DL S004 · Derivatives, the Chain Rule, and Computation Graphs

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 4 / 80 · Module M01 · ~2 hours

The story

Every deep learning framework — PyTorch, JAX, TensorFlow, tinygrad, micrograd, MLX — is, at its core, doing exactly one thing: it is running the chain rule on a big graph, in the right order. That's it. That's the entire magic. Learn the chain rule properly today, and you will never again be afraid of loss.backward(). You will know what those bytes are doing.

Here's the puzzle. Imagine a knob w in a machine somewhere. You turn the knob a tiny bit. Some faraway output L (the loss) changes a tiny bit. How much does L change per unit turn of w? That number is called ∂L/∂w. If you had it, you'd know exactly which direction to turn w to make L smaller — and if you had that number for every knob in the machine, you could improve the machine. That's training. That's the whole game.

The chain rule is what lets us compute ∂L/∂w even when w is buried five layers deep in a mess of matrix multiplications, additions, ReLUs, and softmaxes. It says: to know how a distant output depends on your knob, multiply together all the local sensitivities on the path between them. Elementary calculus, industrial-scale application. By the end of this session you'll have done it by hand on a graph with three ops, then on a mini-MLP with matrices, then in code. When someone tells you PyTorch autograd is "magic", you'll smile.

Prerequisites

• Session 001 — you'll need np.array, shapes, and basic elementwise ops.
• Session 002 — matrix multiply, transpose, shape agreement.
• Session 003 — because gradient shapes must match parameter shapes, and broadcasting is how sums align.

You do not need to have taken calculus recently. We are going to re-derive derivatives from scratch, at 12th-grade speed.

1 · Derivative — what is it, really?

Forget the textbook definition for one moment. Here is the physicist's version:

The derivative of f at point x is the number f'(x) such that if you nudge x by a tiny amount dx, then f(x) changes by approximately f'(x) * dx.

That's it. Sensitivity. Change-per-change. Units matter: if x is in metres and f(x) is in Joules, then f'(x) is in Joules per metre.

Formally:

h is the "nudge". f(x + h) - f(x) is how much f moved. Divide to get "movement per unit nudge". Take the limit so it's exact at the point.

1.1 Numeric check — the finite difference

You can always estimate a derivative numerically. Pick a small h, compute (f(x+h) - f(x)) / h, and that's a decent approximation.

def numgrad(f, x, h=1e-5):
    return (f(x + h) - f(x)) / h
 
# f(x) = x^2, so f'(x) = 2x. At x=3, f'(3) = 6.
f = lambda x: x**2
print(numgrad(f, 3.0))   # 6.00001... ≈ 6.0 ✅

This trick is going to save you multiple times in your career. Every time you write a hand-derived gradient, sanity-check it against a numeric one. Symbolic mistakes are how models silently train wrong.

1.2 Building blocks — memorise these six

You do not need to memorise a calculus textbook. You need six rules. Six.

1.  d/dx (c)              = 0                       constants don't change
2.  d/dx (x^n)            = n * x^(n-1)             power rule
3.  d/dx (e^x)            = e^x                     exponential
4.  d/dx (log x)          = 1 / x                   natural log
5.  d/dx (u + v)          = u' + v'                 sum rule
6.  d/dx (u * v)          = u' * v + u * v'         product rule

Plus one meta-rule (the chain rule), which is §2 and half the point of this whole session.

That's it. Every derivative in this entire 80-session series will be built from these seven things. Print them on a sticky note. Stick it on your monitor. Move on.

1.3 A worked example — one variable

f(x) = 3x² + 2x + 5. What is f'(x)?

Apply sum rule + power rule term-by-term:

d/dx (3x²) = 3 * 2x = 6x
d/dx (2x)  = 2
d/dx (5)   = 0
-----------------
f'(x)      = 6x + 2

Check at x = 1: f'(1) = 8.

Numeric check:

f = lambda x: 3*x**2 + 2*x + 5
h = 1e-5
print((f(1 + h) - f(1 - h)) / (2 * h))   # 7.9999... ≈ 8 ✅

You now can differentiate any polynomial. Good.

2 · The chain rule — the star of the show

Say y = f(g(x)). To turn the outer x knob and know how the outer y changes, you need two sensitivities:

1. how y responds to g (the outer sensitivity, f' evaluated at g(x))
2. how g responds to x (the inner sensitivity, g'(x))

The chain rule says: multiply them.

dy/dx  =  f'(g(x)) * g'(x)

Same rule, in Leibniz notation:

dy/dx  =  dy/dg  *  dg/dx

Notice how the "dg" cancels — that's your mnemonic. It is not literally cancellation (derivatives aren't fractions), but the notation is designed so it looks like it, and you can use that pattern to track long chains.

2.1 Two-step example

y = (3x + 1)². Let g = 3x + 1, so y = g².

dy/dg = 2g            (power rule, treating g as the variable)
dg/dx = 3             (linear)
 
dy/dx = dy/dg * dg/dx = 2g * 3 = 6 * (3x + 1)

Check at x = 1: dy/dx = 6 * 4 = 24.

f = lambda x: (3*x + 1)**2
print((f(1.00001) - f(0.99999)) / (2e-5))   # 23.999... ✅

2.2 Three-step chain — closer to what you'll see in a neural net

y = log(sigmoid(w*x)) where sigmoid(z) = 1 / (1 + e^{-z}).

Break into stages:

z = w * x
s = sigmoid(z) = 1 / (1 + e^{-z})
y = log(s)

Chain rule (backwards, right to left):

dy/ds = 1 / s
ds/dz = s * (1 - s)          (classic — derive it once, use forever)
dz/dw = x
 
dy/dw = dy/ds * ds/dz * dz/dw
      = (1/s) * s * (1-s) * x
      = (1 - s) * x

Neat little cancellation. s disappears. This is not a coincidence — cross-entropy + sigmoid has this simplification baked in, and you'll see the same trick with softmax + cross-entropy in Session 005.

The point: once we have the intermediate values (the "forward pass"), each backward step is a small local calculation. This is why autograd works — you cache the forward values, then walk the graph backwards multiplying local derivatives. That's the whole algorithm.

2.3 The sigmoid-derivative derivation (do this once)

Everyone forgets this the first time. Here it is:

The second factor becomes (1 - s) because e^{-z} / (1 + e^{-z}) = 1 - 1/(1 + e^{-z}) = 1 - s. Beautiful. Write it once on paper.

3 · Partial derivatives — multi-input functions

Neural nets don't take one input; they take vectors and matrices. So we need to extend derivatives to functions of many variables.

A partial derivative is just a normal derivative, computed while pretending every OTHER variable is a constant. Notation: ∂f/∂x (curly-d) instead of df/dx.

Example: f(x, y) = x²y + 3y. Then:

∂f/∂x = 2xy         (treat y as a constant)
∂f/∂y = x² + 3      (treat x as a constant)

Pack them into a vector: ∇f = [∂f/∂x, ∂f/∂y] = [2xy, x² + 3]. That's called the gradient.

The key fact: the gradient points in the direction of steepest ASCENT of f. Negate it → direction of steepest DESCENT. Take a small step in that direction, f decreases. Do this millions of times and you have trained a neural network. That's gradient descent. That's Session 006.

3.1 Chain rule with multiple variables

If y = f(u, v) and both u and v depend on x:

dy/dx = (∂f/∂u) * (du/dx) + (∂f/∂v) * (dv/dx)

Every path from x to y contributes; you sum over paths. This is critical for backprop when a variable feeds into multiple downstream ops.

Worked example. u = x², v = 3x, f(u,v) = u + v — so y = x² + 3x, obviously dy/dx = 2x + 3. Verify with the sum-over-paths rule:

∂f/∂u = 1        du/dx = 2x
∂f/∂v = 1        dv/dx = 3
 
dy/dx = 1*(2x) + 1*(3) = 2x + 3    ✅

Same answer, obtained by walking two paths through the graph and summing. That "sum over paths" is exactly what backprop does when a tensor gets used twice.

4 · Computation graphs — a picture worth 10× the algebra

A computation graph is a directed acyclic graph (DAG) where:

• nodes are values (scalars or tensors)
• edges are ops (add, mul, matmul, ReLU, ...)
• leaves on the left are inputs and parameters (x, w, b)
• the root on the right is the loss L

Every deep learning framework builds one of these graphs (either eagerly, like PyTorch, or ahead-of-time, like old TF). Then backprop is: walk the graph right-to-left, multiplying local derivatives, accumulating at nodes where multiple paths meet.

Let's do it on the simplest interesting example: linear regression on one sample.

4.1 The graph

Expression: L = (w * x + b - y)². Here x and y are data (constant during backprop); w and b are parameters (we want ∂L/∂w and ∂L/∂b).

Let's insert intermediate nodes:

z1 = w * x
z2 = z1 + b
z3 = z2 - y
L  = z3 ** 2

4.2 Forward pass — pick numbers and compute

Let w = 2, x = 3, b = 1, y = 10.

z1 = 2 * 3 = 6
z2 = 6 + 1 = 7
z3 = 7 - 10 = -3
L  = (-3)² = 9

Loss is 9. Big. We want to know how to change w and b to reduce it.

4.3 Backward pass — walk right-to-left, multiply local derivatives

Start at the root with ∂L/∂L = 1. Then, at each node, compute ∂L/∂(this node) using the chain rule.

Node L = z3². Local derivative: ∂L/∂z3 = 2 * z3 = 2 * (-3) = -6.

Node z3 = z2 - y. Local derivative wrt z2: ∂z3/∂z2 = 1. Chain: ∂L/∂z2 = ∂L/∂z3 * ∂z3/∂z2 = -6 * 1 = -6.

Node z2 = z1 + b. Two upstream contributions:

• ∂z2/∂z1 = 1, so ∂L/∂z1 = -6 * 1 = -6.
• ∂z2/∂b = 1, so ∂L/∂b = -6 * 1 = -6.

Node z1 = w * x. Two upstream contributions:

• ∂z1/∂w = x = 3, so ∂L/∂w = ∂L/∂z1 * ∂z1/∂w = -6 * 3 = -18.
• ∂z1/∂x = w = 2 — but we don't care about the gradient of the loss with respect to data.

Answers:

∂L/∂w = -18
∂L/∂b = -6

Interpretation: increasing w by a tiny amount dw decreases L by ~18*dw (since the sign is negative). So to reduce L, we should increase w. That's exactly what gradient descent (Session 006) will do.

4.4 Sanity — numeric check

def f(w, b, x=3.0, y=10.0):
    return (w*x + b - y)**2
 
w, b = 2.0, 1.0
h = 1e-5
gw = (f(w+h, b) - f(w-h, b)) / (2*h)   # ≈ -18.0
gb = (f(w, b+h) - f(w, b-h)) / (2*h)   # ≈ -6.0
print(gw, gb)                          # -18.0000... -6.0000...

✅ Our hand-computed gradients match numerical differentiation. This is what "backprop correctness" means, and it's the check we'll re-run in Sessions 010 and 012.

5 · Coding a mini-autograd — 15 lines of NumPy

Let's crystallize §4 in code. Below is a bare-bones autograd for scalar computation graphs — the ancestor of what Andrej builds in micrograd, which we rebuild in Session 012.

import numpy as np
 
class Value:
    def __init__(self, data, _children=(), _op=""):
        self.data     = float(data)
        self.grad     = 0.0
        self._prev    = set(_children)
        self._op      = _op
        self._backward = lambda: None
 
    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")
        def _backward():
            self.grad  += 1.0 * out.grad
            other.grad += 1.0 * out.grad
        out._backward = _backward
        return out
 
    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), "*")
        def _backward():
            self.grad  += other.data * out.grad
            other.grad += self.data  * out.grad
        out._backward = _backward
        return out

Read those _backward closures carefully. They ARE the chain rule made concrete: "the gradient flowing OUT of this node gets multiplied by my local derivative, and ADDED to the gradient of my input."

Now the neat bit — the reverse-mode driver:

    def backward(self):
        # 1. topological sort: parents before children
        topo, visited = [], set()
        def build(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build(child)
                topo.append(v)
        build(self)
 
        # 2. seed root grad = 1, walk backwards
        self.grad = 1.0
        for v in reversed(topo):
            v._backward()

Why topological sort? Because when you compute ∂L/∂(some node), every downstream contribution must already be accumulated in .grad. Topo order guarantees that. Get this wrong and you compute gradients from stale values.

Try it on our linear regression example:

w = Value(2.0);  x = Value(3.0)
b = Value(1.0);  y = Value(10.0)
 
L = (w*x + b + (-1)*y)     # differences via `+ (-1)*y` to keep ops minimal
L = L * L                   # squared
 
L.backward()
print(w.grad, b.grad)       # -18.0  -6.0   ✅ matches §4.3

Fifteen-ish lines of code and we've built the mathematical core of PyTorch. That's the whole magic. Everything else — tensors, GPUs, matmul, cudnn kernels — is engineering on top of this idea.

6 · Gradients for vector and matrix ops

Real neural nets don't have scalar w * x. They have X @ W + b, where all three are tensors. Backprop still works, but the local derivatives are matrices, and we have to be careful with shapes.

6.1 The linear layer — the atom of every network

Forward pass: Y = X @ W + b.

Shapes: X: (B, D_in), W: (D_in, D_out), b: (D_out,), Y: (B, D_out).

Given some upstream gradient dY (of shape (B, D_out), coming from the layer above), the local derivatives are:

Read the shapes. Each gradient has the SAME shape as the thing it's a gradient of. Memorise that fact — it's the fastest way to catch shape bugs.

6.2 Where do those formulas come from?

Scalar derivation for y_{ij} = Σ_k X_{ik} W_{kj} + b_j:

∂y_{ij} / ∂W_{kj} = X_{ik}
∂y_{ij} / ∂X_{ik} = W_{kj}
∂y_{ij} / ∂b_j    = 1

Sum over the batch and output dims when accumulating dW and db, and you get the matrix formulas above. We'll re-derive this in more detail in Session 010 (backprop by hand) and prove them element-by-element in Session 012.

6.3 The sanity check I always run

Every time I code a new custom layer, I run this:

import numpy as np
np.random.seed(0)
 
B, D_in, D_out = 4, 3, 2
X = np.random.randn(B, D_in)
W = np.random.randn(D_in, D_out)
b = np.random.randn(D_out)
dY = np.random.randn(B, D_out)     # pretend upstream gradient
 
# Analytic
dX = dY @ W.T
dW = X.T @ dY
db = dY.sum(axis=0)
 
# Numeric — check one element of dW as a proof
def loss(W_, X_=X, b_=b):
    Y = X_ @ W_ + b_
    return (Y * dY).sum()          # dot with dY = same as chain rule
 
W_pert = W.copy();  h = 1e-6
W_pert[1, 0] += h;  L_plus  = loss(W_pert)
W_pert[1, 0] -= 2*h; L_minus = loss(W_pert)
print("num  dW[1,0] =", (L_plus - L_minus) / (2*h))
print("anal dW[1,0] =", dW[1, 0])
# Should match to ~1e-8.

If they don't match, your gradient is wrong. Do not proceed until they do. This one check has saved me weeks over the years.

7 · The full mental model — reverse-mode AD in one paragraph

Here it is, the whole thing:

Build the computation graph in the forward pass, caching every intermediate value. Then walk the graph backwards, starting with ∂L/∂L = 1 at the root. At each node, use your cached values to compute your local derivatives (∂output/∂input for each incoming edge), multiply them by the gradient flowing into this node from downstream, and ADD the result to the .grad field of each incoming node. Because we walk in topological order, every node's .grad is fully accumulated before it's read.

That's it. That's reverse-mode automatic differentiation, in one paragraph, no mystery.

Why "reverse mode"? Because we walk edges opposite to the forward direction. There's also forward-mode AD (walk forwards, one input at a time). Reverse mode is dramatically better when you have many inputs (all your weights) and one output (the loss) — which is exactly the shape of every ML problem. That's why every framework uses it.

7b · The 2025 autograd landscape — what production frameworks actually do

The 30-line NumPy autograd in §5 is not a toy. It is, structurally, what PyTorch, JAX, tinygrad, and MLX all do — they just add more ops, more devices, more optimization passes. Let's see how the same chain rule you just implemented scales up to systems that train GPT-4-scale models.

7b.1 · PyTorch — dynamic tape (the industry default)

PyTorch builds the computation graph on the fly during the forward pass, storing each op and its saved tensors in a per-thread tape. When you call .backward(), it walks the tape in reverse. This is called dynamic or define-by-run autograd, and it's why PyTorch code looks like normal Python — there is no separate "compile" step.

The cost: every op has Python overhead, and the graph must be re-built on every forward pass. For small models, this is negligible. For a Llama 3.1 8B training step it's ~10–15% of wall time, which is why PyTorch 2.x introduced torch.compile — a JIT that traces the Python-level graph, fuses ops (via TorchInductor or Triton), and emits a single optimized kernel. Modern training runs (2024–2026) almost always use torch.compile(model) for the 20–40% speedup on the whole training loop.

7b.2 · JAX — pure-functional trace-and-transform

JAX takes the opposite tack. Instead of a tape, you write pure functions and jax.grad(f) returns another pure function that computes the gradient. Under the hood, JAX traces your function once through with abstract inputs, produces an XLA HLO graph, and lets XLA fuse and schedule it.

The genius: jax.jit, jax.grad, jax.vmap, and jax.pmap all compose. jax.jit(jax.grad(jax.vmap(loss))) gives you a compiled, vectorized, differentiated function in one line. This is why almost all of DeepMind's 2024–2025 published training code (Chinchilla, Gemma, Gemini reference implementations) is in JAX.

7b.3 · tinygrad and MLX — the readable frontier

• tinygrad (George Hotz, ~5000 lines): a full autograd engine + graph optimizer + multi-backend runtime, small enough to read in a weekend. If you understood §5, you can read tinygrad's tensor.py and see the same pattern generalized.
• MLX (Apple, 2023–2025): unified memory autograd for Apple Silicon. Same chain rule, tuned for M-series GPUs. Ships with LoRA training scripts that run on a MacBook Air.

Both are worth reading after this session precisely because they are small enough to hold in your head.

7b.4 · Higher-order gradients and second-order optimizers (2024–2025)

Modern research is increasingly hungry for the gradient of the gradient — the Hessian. Uses:

• Shampoo / SOAP (Vyas et al., 2024) — approximate second-order optimizers that beat Adam on LLM pretraining by 30–50% wall time. All the major labs are experimenting with these in 2024–2025.
• Muon (Jordan et al., 2024) — hidden-layer optimizer based on Newton-Schulz iteration. Reportedly used in some Kimi and Moonshot AI training runs.
• Meta-learning — MAML and its descendants need $\nabla_\theta L(\theta - \alpha \nabla_\theta L)$, i.e., a gradient through a gradient.

All of these ride the chain rule you learned today — they just apply it twice. In PyTorch: torch.autograd.grad(loss, params, create_graph=True) gives you a gradient that is itself differentiable.

7b.5 · Gradient checkpointing — trading recomputation for memory

A critical trick you'll meet in every 2024–2025 pretraining recipe. Standard reverse-mode AD stores every intermediate activation from the forward pass so backward can consume it. For a 32-layer transformer with (B, T, d) = (32, 4096, 4096), that's ~64 GB per layer — far more than any GPU has.

Gradient checkpointing (Chen et al., 2016) stores activations only at checkpoint layers, and recomputes intermediate activations on the fly during backward. Cost: ~30% more compute. Benefit: ~$\sqrt{L}$ times less activation memory. Used in every serious LLM training run.

In PyTorch: torch.utils.checkpoint.checkpoint(block, x) wraps a transformer block. In FSDP + HuggingFace: gradient_checkpointing_enable(). Same idea; different API.

Further reading:

• Griewank & Walther, Evaluating Derivatives (2008) — the reference textbook for reverse-mode AD.
• Ansel et al. "PyTorch 2: Faster ML through dynamic Python" (2024).
• Karpathy's micrograd repo — the 100-line teaching autograd. Read it after §5.
• Chen et al. "Training Deep Nets with Sublinear Memory Cost" (2016).

8 · War stories — the gradient bugs that cost me time

db = dY               # shape (B, D_out) — WRONG

db = dY.sum(axis=0)   # shape (D_out,) — RIGHT

grad = (x > 0).astype(float)   # 0 at x=0
# OR
grad = (x >= 0).astype(float)  # 1 at x=0

total_loss = 0
for batch in loader:
    loss = model(batch).loss
    total_loss = total_loss + loss     # BUG: builds a giant graph

9 · Diagram — the two passes side by side

The forward pass carries values. The backward pass carries gradients. Every node is visited once per pass. That's O(N) in the size of the graph, for both passes. That's why deep learning scales — the cost of computing all gradients is the same order as the cost of the forward pass.

This "backward pass is the same cost as forward" fact is one of the great practical miracles of deep learning. Without it, we couldn't train networks with billions of parameters.

10 · Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: Re-code the mini-autograd from §5 tomorrow, from a blank editor. Revisit §7b (autograd landscape) on day 7.

Next session (S005): Probability, entropy, cross-entropy — what number to compute the gradient of. Chain rule (today) + cross-entropy (tomorrow) + gradient descent (S006) = the full basis of every trained neural net.

Sticky note: in code, .grad += ... (never =). In math, backward sums over paths, and every broadcast forward becomes a sum backward.

Recall — from memory

Stretch — for one extra hour

Extend the mini-autograd from §5 with __sub__, __pow__, and a tanh() method. Compute the gradient of tanh(w*x + b) at w=1, x=0.5, b=0.2 and verify against a numeric check to 6 decimal places. Bonus: add exp() and log() and use them to compute the derivative of the softmax + cross-entropy loss on a 3-class example.

In your own words

Explain the chain rule to a friend in one sentence:

Spaced-review pointer

• From S002 — you'll use matrix shapes constantly here; the "gradient shape = parameter shape" rule is a direct consequence of the shape agreement rule for matmul.
• From S003 — every time you broadcast forward, you sum backward. Broadcasting and backprop are dual operations.

Next-session teaser

Now that you can compute how one number depends on another, we need to know what number to compute in the first place — i.e., what makes a good loss function? Session 005 introduces probability, entropy, and cross-entropy. You'll finally understand why "cross-entropy loss" is the same thing as "log-likelihood" and why we log-sum-exp everything in practice. It'll close the loop: chain rule (today) + cross-entropy (S005) + gradient descent (S006) is the full mathematical basis of every neural network you'll ever train.

What to bring back tomorrow — sticky note

• Diagram: the "forward carries values, backward carries gradients" picture from §9.
• Equation: dy/dx = (dy/dg) * (dg/dx) — chain rule.
• Snippet: the _backward closure for *: self.grad += other.data * out.grad. That's the chain rule in code.

Previous: ← DL S003 · Broadcasting · Next: DL S005 · Probability and Information →