6. Automatic differentiation

Why this layer exists

Chapter 5 derived the linear-regression gradient by hand, which took a page for $d + 1$ parameters and one operation between them and the loss. A transformer has millions of parameters and dozens of kinds of operation (matrix multiplies, softmax, normalisation, embedding lookups, masks, residual additions) composed dozens of layers deep, and a change to one layer would mean re-deriving everything downstream of it. Numerical gradients are no way out: they cost two loss evaluations per parameter.

Automatic differentiation makes gradient computation mechanical. Each operation, as it runs, records its inputs and knows the derivative of its output with respect to them at the values it just saw. From that record the chain rule assembles the gradient with respect to every input in one backward sweep, at a small multiple of the cost of the forward computation, however many parameters there are. Nobody derives a whole-model gradient; each operation’s local rule is derived once, in a few lines, and checked once against central differences.

This is what loss.backward() does in PyTorch, and it is what this chapter builds: tinygpt.autograd, first a scalar engine, Value, small enough to hold in your head, then Tensor, the same algorithm on numpy arrays with broadcasting and batched matrix multiplies. Tensor is the engine that trains every model in Part II. The weights in every GGUF file your llama.cpp servers load were produced the same way: a training framework recorded the forward pass and ran reverse-mode autodiff over it, built from local rules like the ones this chapter derives.

Mechanism

Computation graphs

A computation is a sequence of primitive operations. $f = (ab + c)^2$ is three: $u = ab$, then $s = u + c$, then $f = s^2$. Draw every value as a node and connect each operation’s inputs to its output, and the result is a computation graph. It has no cycles, since an operation can only consume values that already exist. The nodes no operation produced, here $a$, $b$ and $c$, are the leaves: inputs, parameters and constants.

The forward pass evaluates the graph in the order it was built: at $a = 2$, $b = -3$, $c = 10$, $u = -6$, $s = 4$, $f = 16$. The backward pass computes the derivative of the output with respect to every node. Write $\bar{v} = \partial f / \partial v$ for the gradient at node $v$. It starts at the output, with $\bar{f} = 1$, and moves one operation at a time toward the leaves, multiplying by each operation’s local derivative:

• $f = s^2$ has $\partial f / \partial s = 2s = 8$, so $\bar{s} = 8$.
• $s = u + c$ has $\partial s / \partial u = \partial s / \partial c = 1$, so $\bar{u} = \bar{s} \cdot 1 = 8$ and $\bar{c} = 8$.
• $u = ab$ has $\partial u / \partial a = b$ and $\partial u / \partial b = a$, so $\bar{a} = \bar{u}\,b = -24$ and $\bar{b} = \bar{u}\,a = 16$.

Differentiating the formula directly agrees: $\partial f / \partial a = 2(ab + c)\,b = 2 \cdot 4 \cdot (-3) = -24$. Each step used only one operation’s local derivative, at the values that operation saw in the forward pass, and the gradient arriving from above. That locality is the whole idea.

The chain rule as a sum over paths

In the example every node fed exactly one operation. In general a node $v$ feeds several, say $w_1, \ldots, w_m$, and $f$ depends on $v$ only through them. Nudge $v$ by a small $\delta$. By chapter 5’s local linear approximation each $w_k$ moves by about $(\partial w_k / \partial v)\,\delta$, and $f$, as a function of the $w$‘s, moves by about the sum of its partials times those moves. The remainders vanish relative to $\delta$ by the argument of chapter 5’s chain rule. Reading off the coefficient of $\delta$,

$\bar{v} = \sum_{k=1}^{m} \bar{w}_k\, \frac{\partial w_k}{\partial v} .$

This is the multivariable chain rule, with one term per edge leaving $v$: an operation that uses $v$ twice, as in $v \cdot v$, contributes two terms. Unrolling the recursion, $\partial f / \partial v$ is the sum, over every path from $v$ to $f$, of the product of the local derivatives along the path. There can be exponentially many paths: a stack of $L$ diamonds, each node feeding two nodes that meet again, has $2^L$. The recursion never enumerates them. Each $\bar{w}_k$ already sums over all paths from $w_k$ onward, so every edge is visited once.

Take $y = x \cdot x + x$ at $x = 3$, with the intermediate $m = x \cdot x$. Then $\bar{y} = 1$ and $\bar{m} = 1$. Three edges leave $x$: one into the addition, local derivative 1, and two into the multiply, where the derivative with respect to each factor is the other factor, 3. So

$\bar{x} = 1 \cdot 1 + 1 \cdot 3 + 1 \cdot 3 = 7 = \frac{d}{dx}\big(x^2 + x\big)\Big|_{x=3} = 2 \cdot 3 + 1 .$

In code, the sum is accumulation: every operation’s backward step adds its contribution into its input’s gradient with +=. With = instead, $\bar{x}$ would be whichever contribution happened to arrive last, 1 or 3.

Order: finish a node before passing it on

A node’s backward step multiplies its gradient $\bar{w}$ by local derivatives and hands the products to its inputs. That is only right if $\bar{w}$ is complete, meaning every consumer of $w$ has already contributed, and if the step runs exactly once. Take the diamond $u = 3a$, $f = u \cdot u + u$ at $a = 2$. The correct values are $\bar{u} = 2u + 1 = 13$ and $\bar{a} = 3 \cdot 13 = 39$. If $u$‘s step ran when only the addition had contributed, $\bar{u} = 1$ and $a$ would receive 3. Running it once more after the multiply’s contribution of 12 arrives would add $3 \cdot 13 = 39$, for a total of 42. Running it only the first time leaves 3. Both are wrong.

The fix is to process every node after all of its consumers. A topological order lists every node after its inputs, as the forward pass created them; reversed, it puts every consumer before its inputs. A depth-first search produces one: visit a node’s inputs first, then append the node.

Forward mode and reverse mode

What the example did is reverse mode: derivatives of the one output with respect to everything, propagated from the output back. The alternative propagates forward. Choose one input $x_i$ and carry alongside every node’s value the number $\dot{v} = \partial v / \partial x_i$, computed as each node is evaluated:

$\dot{w} = \sum_{\text{inputs } v \text{ of } w} \frac{\partial w}{\partial v}\, \dot{v} ,$

seeded with $\dot{x}_i = 1$ and 0 for every other input. At the output, $\dot{f} = \partial f / \partial x_i$. This is forward mode. It multiplies the same local derivatives along the same paths, summed from the other end. One forward-mode pass delivers one partial derivative, so a gradient over $n$ inputs takes $n$ passes (exercise c).

Reverse mode delivers all $n$ partials of one output from one forward pass and one backward pass. The backward pass costs a small constant times the forward pass, because each operation’s backward step does: a multiply does two multiplications where the forward did one, and a matrix multiply, as derived below, does two matrix multiplies of the forward one’s size. A network whose cost is dominated by matrix multiplies computes its gradient for about three times the cost of evaluating it: the forward pass plus a backward pass twice as expensive. Against chapter 5’s numerical gradient, that is three forward passes instead of two million for a single (1000, 1000) weight matrix.

The price is memory: a backward step needs values from the forward pass (the multiply needs $b$ to compute $\bar{a}$), so every intermediate is kept until the backward pass has used it, activation memory that grows with batch size and sequence length as well as depth. Forward mode keeps nothing and wins with few inputs and many outputs. In terms of the Jacobian matrix of all outputs with respect to all inputs, forward mode computes one column per pass and reverse mode one row, and a scalar loss has a Jacobian with a single row.

Gradients with respect to arrays

Tensor has an array per node, and the loss is still a scalar $L$. For a node $\mathbf{X}$ of shape (n, k), the gradient $\bar{\mathbf{X}}$ is the array of partial derivatives $\partial L / \partial X_{ij}$, shape (n, k): the shape of $\mathbf{X}$, as chapter 5’s gradient had the shape of $\theta$. The chain rule applies to every scalar entry unchanged, and each array operation’s rule is that sum worked out once.

Elementwise operations. For $\mathbf{Y} = \exp(\mathbf{X})$, entry $Y_{ij}$ depends on $X_{ij}$ alone, so the sum has one term: $\bar{X}_{ij} = \bar{Y}_{ij}\, e^{X_{ij}} = \bar{Y}_{ij}\, Y_{ij}$. Writing $\odot$ for elementwise multiplication, $\bar{\mathbf{X}} = \bar{\mathbf{Y}} \odot \mathbf{Y}$. The other elementwise functions follow the same pattern with their chapter 5 derivatives: $\bar{\mathbf{Y}} \odot (1 - \mathbf{Y} \odot \mathbf{Y})$ for $\tanh$, $\bar{\mathbf{Y}} / \mathbf{X}$ for $\log$, $\bar{\mathbf{Y}} \odot k\mathbf{X}^{k-1}$ for a constant power $k$, and for $\mathrm{relu}(x) = \max(x, 0)$, $\bar{\mathbf{Y}}$ where $\mathbf{X} > 0$ and 0 elsewhere. For $\mathbf{Z} = \mathbf{X} \odot \mathbf{Y}$ of equal shapes, $\bar{\mathbf{X}} = \bar{\mathbf{Z}} \odot \mathbf{Y}$ and $\bar{\mathbf{Y}} = \bar{\mathbf{Z}} \odot \mathbf{X}$.

The matrix multiply

Let $\mathbf{C} = \mathbf{A}\mathbf{B}$ with $\mathbf{A}$ of shape (n, k), $\mathbf{B}$ of shape (k, m) and $\mathbf{C}$ of shape (n, m), and suppose the backward pass has delivered $\bar{\mathbf{C}}$, shape (n, m). Chapter 3’s index formula is

$C_{ij} = \sum_{p=1}^{k} A_{ip}\, B_{pj} .$

The entry $A_{rs}$ appears in $C_{ij}$ only when $i = r$, as the term $p = s$, with coefficient $B_{sj}$. So $\partial C_{ij} / \partial A_{rs}$ is $B_{sj}$ when $i = r$ and 0 otherwise, and the chain rule sums over the entries $C_{rj}$ for $j = 1, \ldots, m$:

$\bar{A}_{rs} = \sum_{j=1}^{m} \bar{C}_{rj}\, B_{sj} = \sum_{j=1}^{m} \bar{C}_{rj}\, (\mathbf{B}^\top)_{js} = (\bar{\mathbf{C}}\mathbf{B}^\top)_{rs} .$

The middle step rewrites $B_{sj}$ as an entry of the transpose, which turns the sum into chapter 3’s product formula. Shapes: (n, m) times (m, k) is (n, k), the shape of $\mathbf{A}$. For $\mathbf{B}$: $B_{rs}$ appears in $C_{ij}$ only when $j = s$, as the term $p = r$, with coefficient $A_{ir}$, so

$\bar{B}_{rs} = \sum_{i=1}^{n} \bar{C}_{is}\, A_{ir} = \sum_{i=1}^{n} (\mathbf{A}^\top)_{ri}\, \bar{C}_{is} = (\mathbf{A}^\top\bar{\mathbf{C}})_{rs} ,$

shape (k, n) times (n, m), which is (k, m), the shape of $\mathbf{B}$. Together,

$\frac{\partial L}{\partial \mathbf{A}} = \frac{\partial L}{\partial \mathbf{C}}\,\mathbf{B}^\top , \qquad \frac{\partial L}{\partial \mathbf{B}} = \mathbf{A}^\top\,\frac{\partial L}{\partial \mathbf{C}} .$

Each is a matrix multiply with $nkm$ multiply-adds, the same count as the forward product, which is the factor of two in the cost argument above.

Chapter 5’s regression gradient is a special case. Treat chapter 5’s $\mathbf{w}$ as a (d, 1) column (the engine’s @ needs two axes) and drop the bias, so $\hat{\mathbf{y}} = \mathbf{X}\mathbf{w}$; then $L = \frac{1}{n}\sum_i r_i^2$ gives $\partial L / \partial \hat{y}_i = 2r_i / n$. The rule for the right operand gives $\bar{\mathbf{w}} = \mathbf{X}^\top \bar{\hat{\mathbf{y}}} = \frac{2}{n}\mathbf{X}^\top\mathbf{r}$, the formula chapter 5 found by hand.

For stacked matrices, (..., n, k) @ (..., k, m) in chapter 3’s leading-axes sense, every matrix in the stack is an independent product, and the same two rules apply to the last two axes of each: the transpose becomes a swap of the last two axes. A (k, m) weight multiplied against a (B, n, k) batch is broadcast, used once per matrix, and its gradient must add up every use: the broadcasting rule, next.

Broadcasting backward

Chapter 3’s broadcasting repeats an operand along some axes without copying it. Take a bias $\mathbf{b}$, shape (m,), added to $\mathbf{Z}$, shape (n, m): $O_{ij} = Z_{ij} + b_j$. The single number $b_j$ feeds $n$ entries of the output, one per row, each with local derivative 1, so the chain rule sums them:

$\bar{b}_j = \sum_{i=1}^{n} \bar{O}_{ij} ,$

shape (m,). In general, a value repeated along an axis receives the sum of the gradients of its copies along that axis. Chapter 3’s rules say exactly which axes were repeated. Rule 1 padded the shorter shape on the left with size-1 axes, which rule 3 then stretched: those axes do not exist in the operand, so the gradient is summed over them and they are dropped. Rule 3 also stretched any size-1 axis the operand did have: the gradient is summed over it and the axis kept, with size 1. Undoing both, in that order, returns an array of the operand’s shape, and that is _unbroadcast in the walkthrough. For (3, 1) * (1, 4), which broadcasts to (3, 4), the left operand’s gradient is the (3, 4) product gradient summed over axis 1 and kept as (3, 1), and the right operand’s is summed over axis 0, giving (1, 4). Every binary operation in Tensor passes its input gradients through it, so a broadcast anywhere in the forward pass, including a weight shared across a batched matrix multiply, is summed in the backward pass.

Sum, mean and reshaping

Summation runs the other way. For $\mathbf{s} =$ X.sum(axis=1) with $\mathbf{X}$ of shape (n, k), $s_i = \sum_j X_{ij}$, shape (n,), and $\partial s_i / \partial X_{ij} = 1$. Each $X_{ij}$ feeds only $s_i$, so $\bar{X}_{ij} = \bar{s}_i$: the gradient of row $i$‘s sum is copied across row $i$. The backward pass of a sum is a broadcast, as the backward pass of a broadcast is a sum.

The shapes need care. $\bar{\mathbf{s}}$ has shape (n,), and broadcasting it against (n, k) would align it with the last axis, which is the wrong one. The summed axis has to be put back first, as size 1, giving (n, 1), which then broadcasts along the axis that was summed. With keepdims=True the axis was never removed. When $n = k$, skipping this step raises no error: a (k,) gradient broadcasts as a row and gives $\bar{X}_{ij} = \bar{s}_j$, the same trap as chapter 4’s softmax on a square input.

A mean over $N$ entries is a sum times the constant $1/N$, and the engine builds it from those two operations. reshape moves no data, so its backward reshapes the gradient back; a transpose of two axes is undone by swapping them back.

Leaf gradients accumulate; intermediate gradients do not

Both engines give a second call to backward() the same meaning. The gradients of leaves, nodes with no parents, accumulate across calls: after $k$ calls to y.backward(), each leaf’s .grad is $k$ times $\partial y / \partial \text{leaf}$. The gradients of intermediate nodes, anything an operation produced, are reset to zero at the start of each call and recomputed.

Leaves accumulate because the gradient of a sum of losses is the sum of their gradients: calling backward() on the losses of several micro-batches in turn leaves each parameter holding the gradient of their total, so a batch too large for memory is processed in pieces and still takes one step with the full gradient. The cost is a rule every training loop must follow: zero the parameters’ gradients before each step’s backward pass. Chapter 8’s training loop does it, the lab below does it by hand, and the lab shows what happens without it.

Intermediates are reset because otherwise a second call counts the first call again. Suppose nothing were reset, and let $y =$ x.sum() for a leaf x. The first call seeds $\bar{y} = 1$, which passes 1 to each entry of $\bar{\mathbf{x}}$. The second call adds its seed to the stale $\bar{y}$, making 2, and passes 2, so each entry of $\bar{\mathbf{x}}$ is 3 instead of 2. Call $k$ passes $k$, and after $k$ calls the leaf holds

$1 + 2 + \cdots + k = \frac{k(k+1)}{2}$

copies of its gradient instead of $k$: quadratic growth in the number of calls. Each extra level of intermediates raises the degree by one. For $y =$ (x * x).sum(), the multiply’s gradient also accumulates the growing $\bar{y}$, and running the engine with the reset removed gives $\bar{\mathbf{x}}$ equal to 1, 4, 10, 20 and 35 times $2\mathbf{x}$ after one to five calls, $k(k+1)(k+2)/6$, where the answer should be $k$. With the reset, every intermediate gradient is rebuilt from zero, and each call adds exactly one more copy of each leaf’s gradient. test_backward_accumulates pins this: two calls leave $2 \cdot 2\mathbf{x}$.

A constant, the 2.0 in x * 2.0, becomes a leaf too, and accumulates a gradient that nothing reads.

Gradient checking

A wrong backward rule does not crash. It produces a wrong gradient, and training does worse than it should, or diverges, or looks fine on an easy problem. The independent check is chapter 5’s central difference.

gradcheck(f, inputs) runs f and backward(), recomputes each input’s gradient with numerical_grad, and compares the two elementwise by numpy‘s allclose rule,

$\big\lvert g_{\text{autodiff}} - g_{\text{numerical}} \big\rvert \le \text{atol} + \text{rtol} \cdot \lvert g_{\text{numerical}} \rvert ,$

with $\text{atol} = 10^{-6}$ and $\text{rtol} = 10^{-4}$. By chapter 5’s error analysis, with $h = 10^{-6}$ the central difference is good to about $10^{-10}$ when the function and its derivatives are of order 1, and less as $\lvert f \rvert$ grows, since the rounding term scales with it. Over the 18 parametrized gradchecks of the lab tests the largest disagreement between the reference engine and the numerical gradient is $2.0 \times 10^{-8}$, in the batched matrix multiply, whose gradient entries reach 63. The absolute tolerance is 50 times that. A wrong rule, such as a missing factor, a missing transpose or a sum over the wrong axis, is off by amounts comparable to the gradient itself.

Two cases break the comparison rather than the engine. relu has a kink at 0, where it has no derivative; the engine uses 0 there. A central difference that straddles the kink measures something in between: 0.5 at exactly 0, and 0.65 at $3 \times 10^{-7}$, where $(\mathrm{relu}(x + h) - \mathrm{relu}(x - h))/2h = 1.3 \times 10^{-6} / 2 \times 10^{-6}$. So relu is checked at points farther than $h$ from 0; the lab’s case uses $-1.5, -0.3, 0.4, 2.0$. Likewise log and division are checked on inputs in $[0.5, 2]$, away from their pole at 0.

A check at one point is evidence, not proof, and it costs two evaluations of f per input entry, so it runs on small shapes. Those shapes have to exercise every code path: exercise (d) shows a bug that only a square input exposes.

Walkthrough

The reference package is py/tinygpt/autograd/: scalar.py, tensor.py, gradcheck.py, and an __init__.py that exports Value, Tensor, gradcheck and GradcheckError.

The scalar engine

class Value:
    """A scalar that remembers how it was computed, so gradients can flow back through it."""

    def __init__(self, data, _parents=(), _op=""):
        self.data = float(data)
        self.grad = 0.0
        self._parents = _parents
        self._op = _op
        self._backward = lambda: None

    def __repr__(self):
        return f"Value(data={self.data}, grad={self.grad})"

    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 += out.grad
            other.grad += 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

    def __pow__(self, k):
        if not isinstance(k, (int, float)):
            raise TypeError("only constant exponents are supported")
        out = Value(self.data ** k, (self,), f"**{k}")

        def backward():
            self.grad += k * self.data ** (k - 1) * out.grad

        out._backward = backward
        return out

A Value holds its number in data, its gradient in grad, the nodes it was computed from in _parents, and a _backward closure that pushes out.grad into those parents; a leaf’s closure does nothing. The closure reads out.grad when it runs, not when it is defined, so it sees the gradient the backward pass has accumulated by then, and every update is +=. A plain number operand is wrapped in a Value, so a * 2 works; the reflected methods further down the file make 2 * a work too. __pow__ accepts only a constant exponent, since $a^b$ with both variable would also need $\partial / \partial b = a^b \log a$. _op names the operation for debugging.

The rest of the class follows the same pattern. exp reuses its output, since $(e^x)' = e^x$; tanh uses $1 - t^2$ from chapter 5 exercise (b); relu passes the gradient where the input was positive. Negation, subtraction and division are built from multiplication, addition and the power $-1$, and need no rules of their own.

def backward(self):
    """Add d(self)/d(leaf) into .grad of every leaf Value this one depends on."""
    order, seen = [], set()

    def visit(v):
        if id(v) not in seen:
            seen.add(id(v))
            for p in v._parents:
                visit(p)
            order.append(v)

    visit(self)
    for v in order:
        if v._parents:  # intermediate results: recompute from scratch each call
            v.grad = 0.0
    self.grad += 1.0
    for v in reversed(order):
        v._backward()

visit is the depth-first search: it appends a node after all of its parents, so order is a topological order, and seen makes a shared node appear once. Every intermediate node’s gradient is then reset. The seed self.grad += 1.0 is $\partial y / \partial y = 1$; self is normally an intermediate, just reset, so the seed is exactly 1. Finally each _backward runs in reversed order, consumers before their inputs, so every gradient is complete before it is passed on.

The tensor engine

def _unbroadcast(grad, shape):
    """Sum grad down to shape, undoing numpy broadcasting."""
    while grad.ndim > len(shape):
        grad = grad.sum(axis=0)
    for axis, size in enumerate(shape):
        if size == 1 and grad.shape[axis] != 1:
            grad = grad.sum(axis=axis, keepdims=True)
    return grad

_unbroadcast is the two steps of the broadcasting rule. The while loop sums away the leading axes the operand never had. The for loop sums over each axis where the operand has size 1 and the gradient does not, with keepdims=True so the axis stays.

def __add__(self, other):
    other = Tensor._lift(other)
    out = Tensor(self.data + other.data, (self, other), "+")

    def backward():
        self.grad += _unbroadcast(out.grad, self.shape)
        other.grad += _unbroadcast(out.grad, other.shape)

    out._backward = backward
    return out

def __mul__(self, other):
    other = Tensor._lift(other)
    out = Tensor(self.data * other.data, (self, other), "*")

    def backward():
        self.grad += _unbroadcast(out.grad * other.data, self.shape)
        other.grad += _unbroadcast(out.grad * self.data, other.shape)

    out._backward = backward
    return out

def __matmul__(self, other):
    other = Tensor._lift(other)
    if self.ndim < 2 or other.ndim < 2:
        raise ValueError(f"@ needs operands with at least 2 axes, got {self.shape} and {other.shape}")
    out = Tensor(np.matmul(self.data, other.data), (self, other), "@")

    def backward():
        # C = A B  =>  dA = dC B^T,  dB = A^T dC  (on the last two axes)
        self.grad += _unbroadcast(np.matmul(out.grad, np.swapaxes(other.data, -1, -2)), self.shape)
        other.grad += _unbroadcast(np.matmul(np.swapaxes(self.data, -1, -2), out.grad), other.shape)

    out._backward = backward
    return out

def sum(self, axis=None, keepdims=False):
    axes = _axes(axis, self.ndim)
    out = Tensor(self.data.sum(axis=axes, keepdims=keepdims), (self,), "sum")

    def backward():
        g = out.grad if keepdims else np.expand_dims(out.grad, axes)
        self.grad += np.broadcast_to(g, self.shape)

    out._backward = backward
    return out

Tensor has the fields of Value, with data a float64 array and grad a zero array of the same shape. _lift wraps a number or an array in a leaf Tensor. __add__ and __mul__ are the scalar rules with _unbroadcast applied to each input’s gradient. __matmul__ is the derivation, with np.swapaxes(..., -1, -2) as a transpose of every matrix in a stack and _unbroadcast summing the gradient of an operand broadcast across leading axes. It rejects operands with fewer than two axes: np.matmul would promote such a vector to a matrix and drop the added axis from the result, and the backward rule would have to undo that. A caller reshapes a vector to (1, k) or (k, 1) instead, and one rule covers every case.

sum normalises axis with _axes to a tuple of non-negative axis numbers, so None, an int, a negative axis and a tuple take the same path. Its backward step is the mechanism’s: np.expand_dims restores the summed axes as size 1 unless keepdims kept them, and np.broadcast_to stretches the result to the input’s shape. The operations below the listing, the elementwise rules, reshape, transpose, mean and the reflected operators, are written as in the mechanism.

def backward(self, grad=None):
    """Add d(self)/d(x) into x.grad for every leaf tensor x this one depends on."""
    if grad is None:
        if self.data.size != 1:
            raise ValueError("backward() without a gradient needs a single-element tensor")
        grad = np.ones_like(self.data)
    # Iterative post-order DFS: deep graphs must not hit Python's recursion limit.
    order, seen, stack = [], set(), [(self, False)]
    while stack:
        node, expanded = stack.pop()
        if expanded:
            order.append(node)
            continue
        if id(node) in seen:
            continue
        seen.add(id(node))
        stack.append((node, True))
        stack.extend((p, False) for p in node._parents if id(p) not in seen)
    for node in order:
        if node._parents:  # intermediate results: recompute from scratch each call
            node.grad = np.zeros_like(node.data)
    self.grad = self.grad + np.asarray(grad, dtype=np.float64)
    for node in reversed(order):
        node._backward()

For a single-element output the default seed is 1. Any other output needs an explicit seed of its shape, and each leaf then receives the gradient of the scalar $\sum_i \text{grad}_i\, y_i$: test_backward_with_explicit_grad seeds x * 2.0 with $(1, 0, -1)$ and expects $(2, 0, -2)$. Without a seed a multi-element output is an error, since the gradient of a vector has no single meaning.

The depth-first search here is iterative. Value‘s recursive visit uses one Python stack frame per level of the graph, and CPython’s default recursion limit is 1000 frames: Value.backward() on a chain of 20,000 additions raises RecursionError. The graphs of Part II are far shallower, but a graph built in a Python loop reaches the limit easily. The explicit stack holds (node, expanded) pairs. Popping an unexpanded node marks it seen, pushes it back as expanded, then pushes its unseen parents above it. Those parents, and everything above them, are finished before the node’s expanded entry comes back up, so a node is appended to order only after all its parents: the recursion’s post-order. A node pushed twice, by two consumers, is expanded once; the seen check on pop discards the second copy. test_deep_graph_does_not_recurse builds the 20,000-addition chain and requires x.grad to be exactly 1.

Gradient checking

def gradcheck(f, inputs, eps=1e-6, atol=1e-6, rtol=1e-4):
    """Raise GradcheckError unless autodiff gradients of f match numerical ones for every input."""
    arrays = [np.array(x, dtype=np.float64) for x in inputs]
    tensors = [Tensor(a.copy()) for a in arrays]
    out = f(*tensors)
    if not isinstance(out, Tensor) or out.data.size != 1:
        raise ValueError("gradcheck needs f to return a single-element Tensor")
    out.backward()
    for i, t in enumerate(tensors):
        def f_i(x, i=i):
            args = [Tensor(x) if j == i else Tensor(arrays[j]) for j in range(len(arrays))]
            return float(f(*args).data.sum())

        expected = numerical_grad(f_i, arrays[i], eps)
        if not np.allclose(t.grad, expected, atol=atol, rtol=rtol):
            err = float(np.max(np.abs(t.grad - expected)))
            raise GradcheckError(f"input {i}: autodiff and numerical gradients differ (max abs error {err:.3e})")

gradcheck builds fresh leaf tensors from private copies of the inputs, so nothing the caller holds is modified and no stale gradient leaks in, and rejects a function that does not return a single-element Tensor before computing any gradient. f_i is the loss as a function of input i alone, the others held fixed, returning the Python float numerical_grad requires; the default argument i=i binds the current index, where a plain closure would see the loop variable’s final value. GradcheckError names the input and the largest absolute error, and subclasses AssertionError, so a failed check in a test reads as a failed assertion.

Exercises

(a) For $f = \tanh(ab + c)$ at $a = 2$, $b = -1$, $c = 1.5$, compute $\partial f / \partial a$, $\partial f / \partial b$ and $\partial f / \partial c$ by hand, as a backward pass through the graph, and check them with Value.

from tinygpt.autograd import Value

a, b, c = Value(2.0), Value(-1.0), Value(1.5)
f = (a * b + c).tanh()
f.backward()
print(f"f = {f.data:.4f}  df/da = {a.grad:.4f}  df/db = {b.grad:.4f}  df/dc = {c.grad:.4f}")

f = -0.4621  df/da = -0.7864  df/db = 1.5729  df/dc = 0.7864

(b) Derive the backward rule for elementwise division $q = a / b$ from the product and power rules, and check that the graph Tensor.__truediv__ builds computes it.

(c) Show that forward mode needs $n$ passes to compute the gradient of a function of $n$ inputs, and that reverse mode needs one.

(d) Trace the backward step of sum(axis=1, keepdims=False) for an input of shape (3, 4): give the shape of every array involved. What would happen without np.expand_dims, for this input and for a (4, 4) input?

(e) Add a sigmoid operation to Tensor, using the result of chapter 5 exercise (a), and gradcheck it.

    def sigmoid(self):
        s = 1.0 / (1.0 + np.exp(-self.data))
        out = Tensor(s, (self,), "sigmoid")

        def backward():
            self.grad += s * (1.0 - s) * out.grad

        out._backward = backward
        return out

rng = np.random.default_rng(1)
gradcheck(lambda a: (a.sigmoid() ** 2).sum(), [rng.normal(size=(3, 4))])
gradcheck(lambda a: (a.sigmoid() ** 2).sum(), [4.0 * rng.normal(size=(3, 4))])

(f) Implement softmax over the last axis of a Tensor by composing existing operations, and gradcheck it. Then derive the gradient of the cross-entropy loss with respect to the logits directly, and explain why computing it as one fused step, $\mathrm{softmax}(\mathbf{z}) - \mathbf{y}$, is both faster and more numerically stable than differentiating the composed graph.

import numpy as np

from tinygpt.autograd import Tensor, gradcheck


def softmax(z):
    """Softmax over the last axis of z, shape (..., V), composed from Tensor ops."""
    shift = z.data.max(axis=-1, keepdims=True)  # (..., 1) constant; softmax ignores shifts
    e = (z - shift).exp()                        # (..., V)
    return e / e.sum(axis=-1, keepdims=True)     # (..., V) / (..., 1)


def cross_entropy(z, y):
    """Mean over rows of -log softmax(z)[target]; y is one-hot, both (B, V)."""
    return -(softmax(z).log() * y).sum(axis=-1).mean()


rng = np.random.default_rng(1)
w = rng.normal(size=(2, 5))
gradcheck(lambda z: (softmax(z) * w).sum(), [rng.normal(size=(2, 5))])

y = np.eye(6)[[0, 4, 2]]                         # (3, 6) one-hot targets
gradcheck(lambda z: cross_entropy(z, y), [rng.normal(size=(3, 6))])

z = Tensor([[0.0, -800.0]])
with np.errstate(all="ignore"):
    loss = cross_entropy(z, np.array([[0.0, 1.0]]))
    loss.backward()
print(loss.data, z.grad)

inf [[nan nan]]

Lab

Implement Value, then Tensor, in py/labs/ch06/starter.py. gradcheck and GradcheckError are already there, complete: they are the instrument, not the exercise. Run

make lab CH=06 IMPL=mine

until it passes. Read py/labs/ch06/test_lab.py first. The first four tests need only Value: hand-derived gradients, x * x + x giving exactly 7, the unary functions, and the reflected operators; get them passing first, since the algorithm is the same and the scalar version has no shapes to get wrong. Eighteen parametrized cases gradcheck the Tensor operations, including broadcasting, a batched matrix multiply against a shared (3, 4) weight, mean over two axes, and a two-layer tanh network. The rest pin what a gradcheck cannot: gradient shapes (3, 1) and (1, 4) for a (3, 1) * (1, 4) product, accumulation over two calls, ValueError for an unseeded multi-element backward() and for @ on a vector, the explicit seed, the 20,000-node chain, and gradcheck rejecting a no-op backward rule and a multi-element output.

Then train something with your engine. XOR is four points no linear model can fit: the points $(0, 0)$ and $(1, 1)$ map to 0, $(0, 1)$ and $(1, 0)$ to 1, and no line separates the two classes. Chapter 3 showed that stacking linear layers does not help; a nonlinearity between them does. Save this script outside the repository, for example as /tmp/xor_mlp.py:

import numpy as np

from tinygpt.autograd import Tensor

rng = np.random.default_rng(0)
x = Tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]])  # (4, 2)
y = Tensor([[0.0], [1.0], [1.0], [0.0]])                      # (4, 1)

hidden = 4
w1 = Tensor(rng.normal(size=(2, hidden)))                     # (2, 4)
b1 = Tensor(np.zeros(hidden))                                 # (4,)
w2 = Tensor(rng.normal(size=(hidden, 1)) / np.sqrt(hidden))   # (4, 1)
b2 = Tensor(np.zeros(1))                                      # (1,)
params = [w1, b1, w2, b2]

lr = 0.05
for step in range(10_000):
    h = (x @ w1 + b1).tanh()          # (4, 4)
    pred = h @ w2 + b2                # (4, 1)
    loss = ((pred - y) ** 2).mean()   # ()
    for p in params:
        p.grad = np.zeros_like(p.data)
    loss.backward()
    for p in params:
        p.data -= lr * p.grad
    if step % 100 == 0 or loss.data < 0.01:
        print(f"step {step:4d}  loss {loss.data:.5f}")
    if loss.data < 0.01:
        break

print("predictions", pred.data[:, 0].round(3))

and run it from py/ with PYTHONPATH=. uv run python /tmp/xor_mlp.py. Change the import to from labs.ch06.starter import Tensor to run yours. The reference engine prints:

step    0  loss 0.76576
step  100  loss 0.22574
step  200  loss 0.18359
step  300  loss 0.13689
step  400  loss 0.08256
step  500  loss 0.03535
step  600  loss 0.01078
step  606  loss 0.00995
predictions [0.07  0.904 0.919 0.138]

A correct engine of yours should agree to the printed digits, or differ in the last one where a different but valid order of floating-point operations rounds differently.

• The model is chapter 3’s affine layer twice with a tanh between, written in Tensor operations: (4, 2) @ (2, 4), a (4,) bias broadcast over the four examples (and its gradient summed back by _unbroadcast), then (4, 4) @ (4, 1). The second weight is scaled by chapter 3’s $1/\sqrt{\text{in}}$ so the initial outputs are of order 1.
• The step is chapter 5’s gradient descent with the gradient from the engine: forward, zero every parameter’s gradient, backward(), move each parameter against its gradient. The printed predictions are the ones the final loss was computed from.
• The curve does not fall by a constant factor per step, as chapter 5’s quadratic did: 0.766 to 0.226 in the first hundred steps, only to 0.184 in the next hundred, then 0.137 to 0.036 over steps 300 to 500. A network with a hidden layer has no single curvature; chapter 5’s analysis applies only locally.

Finally, delete the two lines that zero the gradients and run it again. Each step then subtracts the learning rate times the sum of every gradient computed so far. The loss reaches 0.023 at step 157, lower than the correct run gets before step 600, then overshoots and never recovers: from step 1,000 to step 10,000 it stays between 0.167 and 0.192, and the final predictions are $(0.274, 0.624, 0.624, 0.624)$. That is leaf-gradient accumulation doing what it was designed to do, in a loop that did not ask for it.

Further reading

• Atılım Güneş Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul and Jeffrey Mark Siskind, “Automatic Differentiation in Machine Learning: a Survey”, Journal of Machine Learning Research (2018). Forward and reverse mode, how autodiff differs from numerical and symbolic differentiation, and how frameworks implement it.
• Andreas Griewank and Andrea Walther, Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation (2nd ed., 2008), chapters 1–3. The standard reference: evaluation procedures as computation graphs, and forward and reverse mode, stated precisely.
• David E. Rumelhart, Geoffrey E. Hinton and Ronald J. Williams, “Learning Representations by Back-propagating Errors”, Nature (1986). The chain rule run backward through a layered network to train hidden units; this chapter’s backward pass is the same computation, organised around a graph instead of layers.
• Andrej Karpathy, micrograd (GitHub, 2020). A scalar reverse-mode engine and a small neural-network library on top of it, about 100 and 50 lines of Python. The scalar engine here follows its design.