Derivatives are local measurements; backprop is bookkeeping

Course map · DL004 · Calculus and automatic differentiation

The story

A parameter is a knob; the loss is a measurement downstream. The derivative asks how much that measurement changes locally per unit movement of the knob. The chain rule relates the knobs even when several operations separate them. This is the mathematical core of backpropagation, not a claim that an entire framework is a tiny scalar class.

We will derive a derivative from a limit, multiply sensitivities along a chain, add contributions at a branch, and implement a reverse traversal. Then we will test matrix gradients, higher derivatives and checkpointing on CPU. None of these experiments trains a language model.

Prerequisites

Use DL001–003's arrays, matrix products, singleton axes and reductions. A gradient of a scalar with respect to a tensor has that tensor's shape. Keep that rule beside the page, but remember that shape agreement alone cannot prove a derivative correct.

1 · Derivative — what is it, really?

For a differentiable scalar function,

The derivative's units are output units divided by input units: joules per metre, for example. It is a local exchange rate, like how many extra lumens a dimmer produces per degree at its present setting. A positive derivative says a sufficiently small negative perturbation lowers a smooth function when the derivative is nonzero. It does not authorize an arbitrarily large step.

1.1 Numeric check — the finite difference

For $f(x)=x^2$, the forward quotient is $2x+h$, so its truncation error is exactly $h$. The centred quotient cancels this term. More generally, Taylor expansion gives a centred error of order $h^2$ for a sufficiently smooth function. Roundoff introduces an opposing error that grows roughly as machine precision divided by $h$, with scale-dependent constants. There is no universally optimal 1e-5.

For $x^3-2x$ at $x=2$, the derivative is $3(2)^2-2=10$ and the exact centred quotient is $10+h^2$. A tiny step eventually subtracts nearly equal floating-point values and can get worse. Hold random inputs and masks fixed; test several step sizes. At a kink such as absolute value or ReLU at zero, a centred difference can average two slopes without being a classical derivative.

1.2 Building blocks — memorise these six

A constant differentiates to zero; $x^n$ to $nx^{n-1}$ on its valid domain; $e^x$ to itself; natural log to $1/x$ for positive real $x$; a sum to the sum of derivatives; a product $uv$ to $u'v+uv'$. The chain rule composes these rules. Domains matter: an invalid logarithm in the forward pass is not repaired by a formally correct backward formula.

1.3 A worked example — one variable

For $f(x)=3x^2+2x+5$, the three terms contribute $6x$, $2$ and $0$. Thus $f'(1)=8$. The constant changes the function value but not its slope. This distinction explains why subtracting a loss value from a weight is not gradient descent: the loss and its sensitivity are different quantities with different units.

2 · The chain rule — the star of the show

For $y=f(g(x))$, $dy/dx=f'(g(x))g'(x)$. Evaluate the outer derivative at the inner value, not at the original input. Leibniz notation $(dy/dg)(dg/dx)$ is a useful mnemonic; it is not permission to cancel arbitrary partial-derivative symbols as fractions.

2.1 Two-step example

Let $g=3x+1$ and $y=g^2$. At $x=1$, $g=4$, the outer sensitivity is 8 and the inner sensitivity is 3. Their product is 24. The complete derivative is $6(3x+1)$, not $6x$.

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

For $z=wx$, $s=\sigma(z)$ and $y=\log s$, the successive derivatives are $1/s$, $s(1-s)$ and $x$. Therefore $dy/dw=(1-s)x$. This is a log-probability derivative; a negative log-likelihood reverses its sign. At $w=0,x=2$, $s=1/2$, so the log-probability derivative is 1.

2.3 The sigmoid-derivative derivation (do this once)

Differentiate $(1+e^{-z})^{-1}$: the outer reciprocal gives $-(1+e^{-z})^{-2}$ and the inner derivative gives $-e^{-z}$. The negatives cancel. Factor the result as

The maximum is $1/4$ at $s=1/2$. That is a bound on this local activation derivative, not the full derivative through arbitrary weight matrices.

3 · Partial derivatives — multi-input functions

For $f(x,y)=x^2y+3y$, hold the other independent coordinate fixed: $\partial_x f=2xy$, $\partial_y f=x^2+3$. At $(2,1)$ the gradient is $(4,7)$. For a unit Euclidean direction $v$, the directional derivative is $\nabla f\cdot v$; Cauchy–Schwarz makes the normalized gradient steepest ascent in that norm. Other norms or parameterizations change the steepest direction.

3.1 Chain rule with multiple variables

If $u=x^2$, $v=3x$ and $y=u+v$, there are two paths from $x$ to $y$. Add their contributions: $1(2x)+1(3)=2x+3$. A shared value must receive all downstream contributions. The gradient is local, not a compass pointing at the global minimum: at $(0.1,10)$ on $100x^2+y^2$, the gradient is $(20,20)$ while the direction to the origin is $(-0.1,-10)$. The normalized descent direction has cosine about 0.714 with the direction to the minimum, not 1.

4 · Computation graphs — a picture worth the algebra

A graph records values and their dependencies. One may draw operations as nodes or as labeled edges; be consistent. Reverse traversal accumulates sensitivities after all downstream uses have contributed.

4.1 The graph

For $L=(wx+b-y)^2$, introduce $z_1=wx$, $z_2=z_1+b$, $z_3=z_2-y$, $L=z_3^2$. Data need not be optimized, but derivatives with respect to data still exist and are useful for sensitivity analysis.

4.2 Forward pass — pick numbers and compute

With $w=2,x=3,b=1,y=10$, the values are $z_1=6,z_2=7,z_3=-3,L=9$. Write both values and sensitivities beside the nodes; confusing an activation with its gradient is a common implementation error.

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

Seed $dL/dL=1$. Squaring gives $dL/dz_3=-6$. Subtraction passes -6 to $z_2$ and +6 to the target. Addition passes -6 to $z_1$ and -6 to $b$. Multiplication passes $-6x=-18$ to $w$ and $-6w=-12$ to $x$. Thus increasing $w$ slightly lowers the loss. For a parameter update, use the old forward values until the whole backward pass finishes.

4.4 Sanity — numeric check

A second, intentionally different example uses the half squared loss: at $x=2,y=1,w=3,b=-1$, $a=5,L=8$, and gradients are 8 and 4. A step of 0.1 gives $w'=2.2,b'=-1.4,a'=3,L'=2$. The factor of one-half accounts for the differing gradient formula.

import math
 
def f(w):
    return .5 * (2*w - 2)**2
 
for h in (1e-3, 1e-4, 1e-5):
    numeric = (f(3+h) - f(3-h)) / (2*h)
    assert math.isclose(numeric, 8., abs_tol=1e-7)
print('central differences agree with 8')

The exact retained numerical witness checks that second example. The supplemental program below also checks the original $-18,-6$ example and the derivative-building exercises. Agreement is evidence at the selected inputs, not a proof over every branch.

5 · Coding a mini-autograd — a complete scalar engine

Let $u=w^2,L=u+u$. At $w=3$, $L=18$ and $dL/dw=12$. Deduplicate nodes in the topological traversal, but never deduplicate their uses. w*w has two contributing operand positions, and u+u has two contributing addends.

class Value:
    def __init__(self, data, parents=()):
        self.data, self.grad = float(data), 0.0
        self.parents = parents
        self.pullback = lambda: None
 
    def __add__(self, other):
        out = Value(self.data + other.data, (self, other))
        def backward():
            self.grad += out.grad
            other.grad += out.grad
        out.pullback = backward
        return out
 
    def __mul__(self, other):
        out = Value(self.data * other.data, (self, other))
        def backward():
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out.pullback = backward
        return out
 
    def backward(self):
        order, seen = [], set()
        def visit(node):
            if node not in seen:
                seen.add(node)
                for parent in node.parents:
                    visit(parent)
                order.append(node)
        visit(self)
        for node in order:
            node.grad = 0.0
        self.grad = 1.0
        for node in reversed(order):
            node.pullback()
 
w = Value(3)
u = w * w
loss = u + u
loss.backward()
assert loss.data == 18 and w.grad == 12
print('shared-path gradient', w.grad)

This exact retained engine deliberately accepts Value operands rather than arbitrary numeric coercion, and resets every reachable gradient for each backward call. That makes replay on an unchanged graph idempotent. PyTorch instead accumulates into leaf .grad fields across backward calls and normally releases saved graph data. Do not use one contract to explain the other. The old example (-1)*y required a reflected multiplication method it did not define; Value(-1)*y works with this engine.

Do not mutate Value.data between the forward pass and any backward replay. Multiplication reads its parents’ data in the pullback; this teaching engine has no mutation/version checks. Rebuild the forward graph after changing a parameter. Resetting gradients makes replay idempotent only while those forward values stay unchanged.

The driver uses reverse topological order, not merely reversing construction timestamps. A closure applies a local vector–Jacobian product. This engine is scalar and recursive, so it is not a large-graph implementation, compiler, device runtime or framework-equivalent replica.

6 · Gradients for vector and matrix ops

6.1 The linear layer — the atom of many networks

For $Y=XW+b$, shapes are $(B,D)$, $(D,H)$, $(H,)$ and $(B,H)$. Given an arbitrary upstream cotangent $G$ of shape $(B,H)$,

6.2 Where do those formulas come from?

Expand $Y_{ij}=\sum_kX_{ik}W_{kj}+b_j$ and $dL=\sum_{ij}G_{ij}dY_{ij}$. The coefficient of $dW_{kj}$ is $\sum_iX_{ik}G_{ij}$; the coefficient of $dX_{ik}$ is $\sum_jG_{ij}W_{kj}$; and the coefficient of $db_j$ is $\sum_iG_{ij}$. This derivation also explains why a broadcast forward sums backward. For more leading axes, sum over all replicated axes and restore the original singleton dimensions.

A batch mean divides once, normally inside $G$. Do not average dW again if $G$ already includes the denominator. A shape-correct gradient can still be wrong by a factor of batch size.

6.3 The sanity check I always run

Use the scalar witness $L=\sum Y\odot G$. It makes arbitrary upstream gradients testable without building the rest of a network. Check every coordinate of small X, W and b, not only one conveniently chosen weight. Reject a deliberately wrong db=G.mean(0) to ensure the test detects a realistic bug.

import numpy as np
import math
 
def central(f, x, h=1e-5):
    return (f(x+h)-f(x-h))/(2*h)
 
assert abs(central(lambda x: x**3-2*x, 2.)-10) < 1e-7
assert abs(central(lambda x: 3*x*x+2*x+5, 1.)-8) < 1e-7
assert abs(central(lambda x: (3*x+1)**2, 1.)-24) < 1e-7
assert abs(central(lambda w: (3*w+1-10)**2, 2.)+18) < 1e-7
assert abs(central(lambda b: (6+b-10)**2, 1.)+6) < 1e-7
w = Value(2.); b = Value(1.)
e = w*Value(3.) + b + Value(-10.)
L = e*e; L.backward()
assert (L.data, w.grad, b.grad) == (9., -18., -6.)
L.backward(); assert w.grad == -18.  # this engine resets, not accumulates
 
# Extend without changing the protected Value implementation.
def neg(x): return x*Value(-1.)
def sub(x, y): return x+neg(y)
def power(x, n):
    assert isinstance(n, int) and n >= 0
    out = Value(1.)
    for _ in range(n): out = out*x
    return out
 
def unary(x, fn, derivative):
    value = fn(x.data)
    out = Value(value, (x,))
    slope = derivative(x.data, value)  # save forward-time derivative
    def back(): x.grad += slope*out.grad
    out.pullback = back
    return out
 
def vtanh(x): return unary(x, math.tanh, lambda z, a: 1-a*a)
def vexp(x): return unary(x, math.exp, lambda z, a: a)
def vlog(x):
    assert x.data > 0
    return unary(x, math.log, lambda z, a: 1/z)
 
w = Value(1.); b = Value(.2)
a = vtanh(w*Value(.5)+b); a.backward()
assert abs(w.grad-.5*(1-math.tanh(.7)**2)) < 1e-12
assert abs(w.grad-central(lambda w: math.tanh(.5*w+.2), 1.)) < 1e-8
x = Value(3.); power(sub(x, Value(1.)), 3).backward()
assert x.grad == 12.
z = [Value(2.), Value(1.), Value(.1)]
# Fixed shift equals current max; no differentiation through its selection needed.
e = [vexp(sub(t, Value(2.))) for t in z]
loss = sub(vlog(e[0]+e[1]+e[2]), sub(z[0], Value(2.)))
loss.backward()
p = np.exp([0., -1., -1.9]); p /= p.sum(); p[0] -= 1
np.testing.assert_allclose([t.grad for t in z], p, atol=1e-12)
 
rng = np.random.default_rng(4)
X = rng.normal(size=(4,3)); W = rng.normal(size=(3,2))
b = rng.normal(size=2); G = rng.normal(size=(4,2))
args = [X, W, b]; analytic = [G@W.T, X.T@G, G.sum(0)]
def objective(): return float(((X@W+b)*G).sum())
for arr, expected in zip(args, analytic):
    numeric = np.zeros_like(arr)
    for idx in np.ndindex(arr.shape):
        old = arr[idx]; arr[idx] = old+1e-5; plus = objective()
        arr[idx] = old-1e-5; minus = objective(); arr[idx] = old
        numeric[idx] = (plus-minus)/2e-5
    np.testing.assert_allclose(numeric, expected, atol=1e-8)
assert not np.allclose(G.mean(0), analytic[2])
assert central(lambda x: max(0., x), 0.) == .5  # not our zero subgradient
print('scalar extensions, all matrix gradients and negative controls passed')

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

Evaluate the forward graph, saving values needed by each local backward rule. Seed an output cotangent, then traverse dependencies in reverse topological order. Multiply the accumulated cotangent by each local Jacobian without normally materializing that Jacobian, and add every contribution to the input cotangents. A scalar loss needs one reverse sweep for all parameter gradients. Forward mode propagates an input direction, which is preferable for many outputs and few input directions; neither mode is universally superior.

For ordinary differentiable primitives with inexpensive pullbacks, backward has work of the same asymptotic order as forward, often a small constant multiple. Tensor operation cost matters: counting graph nodes alone treats a scalar addition and a large matrix multiply as equal. Memory depends on what each rule saves, not necessarily every forward intermediate. Checkpointing changes that tradeoff.

7b · Autograd implementations and higher derivatives

7b.1 · PyTorch — dynamic dependency graph

PyTorch records autograd operations during the forward pass and reconstructs the graph on subsequent passes, allowing ordinary executed Python control flow. Saved tensors support backward. model.eval() changes selected module behaviour, not gradient recording; no_grad, inference mode and detach are distinct controls. A detached tensor can share storage with its source.

torch.compile uses graph capture, AOTAutograd and compiler backends; it need not produce one kernel and can encounter graph breaks or recompilations. The old 10–15% overhead and universal 20–40% speedup were unsupported. This lesson tests eager CPU autograd only, not compilation throughput.

7b.2 · JAX — function transformations

JAX's grad transforms a scalar-output function; vmap batches functions and jit compiles supported computations. grad(vmap(loss)) is generally invalid because batching makes the output nonscalar. Use vmap(grad(single_loss)) for per-example gradients, or grad(lambda w: mean(vmap(single_loss)(w, batch))) with appropriate mapped axes for an average loss. Pure-function discipline and tracing constraints are part of the API, not proof that arbitrary stateful Python can be differentiated. No JAX installation is needed for this source-reviewed comparison.

7b.3 · tinygrad and MLX — reading another implementation

tinygrad and Apple's MLX offer other implementations of tensor differentiation. MLX documents NumPy-like APIs, function transformations, lazy evaluation and shared CPU/GPU memory on Apple silicon. These are implementation choices around differentiation, not evidence that either library is a fixed number of lines or readable in a weekend. To study one, trace an add/multiply backward rule, a shared-value test, gradient reset semantics and a device synchronization point. We do not run either framework here.

7b.4 · Higher-order gradients and optimizer distinctions

For $f(\theta)$, the Hessian contains derivatives of the gradient. Often one needs only $Hv$, a Hessian-vector product, not a dense parameter-squared matrix. Differentiating $\nabla f\cdot v$ computes it. MAML-style objectives differentiate through a parameter update and can require these terms.

Shampoo builds matrix preconditioners from gradient statistics; SOAP uses Adam in a Shampoo-derived eigenbasis. Muon transforms momentum matrices using approximate orthogonalization. These are not all “chain rule twice” or explicit Hessian methods. Sophia includes curvature-estimation variants; a name associated with second-order optimization does not identify the estimator or guarantee a speedup. DL006 derives their distinctions.

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

For a simple chain of L equal-sized layers, saving every k-th boundary retains about L/k checkpoints; recomputing one segment needs about k intermediate activations. The toy peak model is $M(k)=A(L/k+k)$, minimized at $k=\sqrt L$. For L=100, k=10 gives about20 activation units rather than100; this ignores parameters, gradients and workspaces. It is not a universal square-root saving or 30% overhead law.

The original (32,4096,4096) activation holds 536,870,912 elements: 1 GiB at two bytes each, not64 GB per activation. A layer can retain several such tensors, so this is not a whole-layer memory estimate. Checkpointed re-execution must preserve the relevant computation, including randomness and state. PyTorch2.9 recommends explicitly choosing use_reentrant=False; API variants differ.

import torch
from torch.utils.checkpoint import checkpoint
 
theta = torch.tensor([2., -1.], dtype=torch.float64, requires_grad=True)
v = torch.tensor([.3, .7], dtype=torch.float64)
f = (theta**4).sum()
g, = torch.autograd.grad(f, theta, create_graph=True)
hv, = torch.autograd.grad((g*v).sum(), theta)
torch.testing.assert_close(hv, 12*theta.detach()**2*v)
 
def block(x): return torch.tanh(x@matrix)
matrix = torch.tensor([[1., .3], [-.2, .8]], dtype=torch.float64)
a = torch.tensor([[.5, -.7]], dtype=torch.float64, requires_grad=True)
b = a.detach().clone().requires_grad_()
y1 = block(a); y2 = checkpoint(block, b, use_reentrant=False)
y1.sum().backward(); y2.sum().backward()
torch.testing.assert_close(y1, y2); torch.testing.assert_close(a.grad, b.grad)
assert 32*4096*4096*2 == 1024**3
print('CPU higher-order and checkpoint output/gradient contracts passed')

8 · Gradient failure cases and method tradeoffs

• Bias not reduced: assigning db=G leaves an extra batch axis; averaging instead of summing can silently rescale a valid shape. Derive the reduction from sharing.
• Accumulation versus reuse: intentional microbatch accumulation differs from stale gradients. Zero leaves between independent PyTorch updates; our scalar engine resets the entire reachable graph. Repeated backward on a retained framework graph has additional lifetime rules.
• ReLU at zero: zero is common under padding, initialization and quantization—not “measure zero on floats.” Record the selected derivative convention; a finite-difference disagreement at the kink is expected.
• Logging retains graphs: storing total_loss = total_loss + loss can retain dependencies. For scalar logging use a detached Python value such as loss.item(); do not detach the loss that should train the model. GPU .item() can also synchronize, a separate cost.
• Nonfinite values: differentiating an invalid forward computation can propagate NaNs. Inspect the first failing operation rather than blaming the last gradient.

Finite differences work with deterministic black-box evaluations but cost two evaluations per coordinate for the centred version. Forward differences can reuse a baseline and need one extra per coordinate. Analytic formulas are inspectable and sometimes faster but need maintenance; autograd composes registered local rules, yet custom incorrect rules, nondifferentiable branches and numerical instability remain possible. Pair an analytic or autograd implementation with independent numerical and algebraic witnesses rather than calling any method infallible.

The forward pass carries values; the backward pass carries cotangents. At a forward split, backward contributions add. A residual path contributes an identity term, but does not guarantee well-conditioned gradients. Products of Jacobians can vanish or explode; inspect the actual weights, activation derivatives and normalization, not just the number of layers.

10 · Retention scaffold

1. State both chain-rule notations. $f'(g(x))g'(x)$ and $(dy/dg)(dg/dx)$ describe the same local composition. For $(3x+1)^2$ at1 they give $8\times3=24$.
2. Re-derive sigmoid's derivative. The reciprocal contributes a minus sign and $e^{-z}$ contributes another; factoring leaves $s(1-s)$, whose peak is1/4.
3. Diagnose two backward bugs. Overwriting shared contributions reports6 instead of12 for $w^2+w^2$ at3; failing to sum a broadcast bias leaves (B,H) instead of (H,). A mean instead of sum is the same shape but wrong magnitude.
4. Give all linear gradients and shapes. $GW^T$ is (B,D), $X^TG$ is (D,H), and G.sum(0) is (H,). A loss denominator belongs exactly once.
5. When is forward mode attractive? For one input direction and many outputs, a single forward directional sweep can deliver all output sensitivities. Reverse mode needs one seed per output for the full Jacobian; compare dimensions before choosing.
6. Derive the checkpoint rule. Differentiate L/k+k to obtain $-L/k^2+1=0$, hence k≈√L; round to feasible segment lengths and measure real memory. Saved RNG state, workspaces and uneven layers break the simple cost model.

Recall — from memory

Explain why graph nodes are visited once but edges can contribute more than once; reconstruct the values9 and gradients-18,-6; and distinguish a scalar backward() seed from a supplied cotangent for a nonscalar output. “Gradient shape equals parameter shape” is a consequence of differentiating a scalar with respect to each coordinate, not a complete test.

Stretch — for one extra hour

The tested extension implements subtraction, nonnegative integer powers, tanh, exp and positive-domain log around the retained engine. Rebuild it without reading. At $\tanh(.5w+.2)$ and w=1, the weight derivative is $0.5(1-\tanh(.7)^2)$, about0.31737; the bias derivative is twice that. The three-class log-softmax example produces softmax minus one-hot. Test a reused input to every new binary operation and a second backward call; a test without shared paths misses the engine's central difficulty.

In your own words

Explain how local sensitivities can compute all gradients without constructing a full Jacobian. Include “multiply along chains, add at shared uses.”

Spaced-review pointer

Tomorrow, rebuild the shared-value graph and all-coordinate matrix check. A week later, derive the checkpoint cost model and explain why a compiler optimization is a different claim from derivative correctness.

Next-session teaser

DL005 chooses the loss we differentiate: probability, entropy, cross-entropy and stable logit calculations.

What to bring back tomorrow — sticky note

Use old forward values, accumulate every edge, unbroadcast gradients, and verify at more than one shape and step size.

Primary references and boundary

PyTorch2.9 autograd, checkpoint, JAX differentiation, MLX0.29.0, and Chen et al.2016 support the implementation comparisons. Source review and local CPU fixtures are not JAX/MLX/GPU execution, compiler benchmarks or production-training evidence.