Chapter 10

Automatic Differentiation

Computational graphs, reverse mode, vector-Jacobian products, and a tiny NumPy autograd.

Automatic differentiation is the machinery that turns a loss into gradients for millions or billions of parameters. It is not symbolic algebra and not numerical finite differences; it is ordinary program execution plus the chain rule applied to every primitive operation. That difference matters: finite differences scale with the number of parameters and suffer from step-size error, while autodiff gives machine-precision derivatives for the executed program. For LLMs in 2026, autodiff is the reason a Python model definition can become a training step.

10.1 Computational graphs

A computation can be viewed as a directed acyclic graph. Leaves hold inputs and parameters; interior nodes hold primitive operations such as add, multiply, matrix multiply, exp, log, relu, and sum; the final node is often a scalar loss. During the forward pass, each node stores its value. During the backward pass, each node receives an adjoint, the derivative of the final loss with respect to that node’s value.

For a node $\vv = f(\vu)$ , the chain rule says changes in $\vu$ affect the loss through $\vv$ :

\frac{\partial L}{\partial \vu} = \frac{\partial L}{\partial \vv}\frac{\partial \vv}{\partial \vu} .\tag{10.1}

The graph matters because one value can feed several later operations. Reverse mode must add all contributions to that value’s adjoint. The tests for this chapter include a shared scalar used twice, precisely to catch missing accumulation. Finite differences would estimate the same derivatives by rerunning the whole program under tiny perturbations; autodiff instead reuses the exact intermediate values from the forward pass and applies analytic local rules.

10.2 Forward mode and reverse mode

Forward-mode autodiff pushes a tangent alongside each value. If $\dot{\vu}$ is a small input perturbation, the local rule is a Jacobian-vector product:

\dot{\vv} = \mJ_f(\vu)\dot{\vu} .\tag{10.2}

One forward sweep gives the derivative of every output in one chosen input direction. That is excellent when there are few inputs and many outputs. It is also useful for checking one directional derivative without storing a whole backward graph, but it is a poor default when the input vector is the full parameter set of a neural network.

Reverse mode runs the primal computation first, then walks the graph backward. Each operation knows how to turn an output adjoint into input adjoints:

\bar{\vu} = \bar{\vv}\,\mJ_f(\vu) .\tag{10.3}

This vector-Jacobian product is why reverse mode wins for deep learning. The full Jacobian is usually never materialized; each primitive consumes an incoming vector and emits vectors for its parents. Training usually has one scalar loss and many parameters. One reverse sweep computes the gradient of that scalar with respect to every parameter, instead of one forward sweep per parameter.

10.3 Vector-Jacobian products

A backward rule is a local VJP. For $z=x+y$ , the output adjoint flows unchanged to both inputs. For $z=xy$ , the rules are $\bar{x} \mathrel{\char"2B}= \bar{z}y$ and $\bar{y} \mathrel{\char"2B}= \bar{z}x$ . Broadcasting adds one wrinkle: if $y$ was stretched across rows, its adjoint must be summed back to the original shape. This is the adjoint of copying: if one parameter value influenced many outputs, all those output sensitivities must be added before updating that parameter. The shared unbroadcast helper from Section B.2 does that reduction.

For matrices, the VJP of $\mY = \mA\mW$ is the familiar pair

\bar{\mA} = \bar{\mY}\mW^\T, \qquad \bar{\mW} = \mA^\T\bar{\mY} .\tag{10.4}

Those formulas are the same matrix calculus used in Section 9.2; autodiff just applies them mechanically across the whole graph. This locality is what makes new layers manageable: implement a correct forward computation and a VJP for its inputs, then the engine composes it with every other layer.

10.4 A tiny reverse-mode Tensor

The Tensor below stores a NumPy value, a gradient buffer, its parents, and a closure that implements the local backward rule. Addition and multiplication demonstrate adjoint accumulation and broadcasting-aware gradient reduction.

Listing 10.1 A minimal Tensor core with broadcasting-aware VJPs

class Tensor:
    def __init__(self, data, _children=()):
        self.data = np.asarray(data, dtype=np.float64)
        self.grad = np.zeros_like(self.data)
        self._prev = tuple(_children)
        self._backward = lambda: None

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

        def _backward():
            self.grad += unbroadcast(out.grad, self.data.shape)
            other.grad += unbroadcast(out.grad, other.data.shape)
        out._backward = _backward
        return out

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

        def _backward():
            self.grad += unbroadcast(out.grad * other.data, self.data.shape)
            other.grad += unbroadcast(out.grad * self.data, other.data.shape)
        out._backward = _backward
        return out

More operations are just more local VJPs. The sum rule broadcasts the output adjoint back to the input shape. The exp, log, and relu rules use their elementary derivatives.

Listing 10.2 Matrix multiply, reductions, and elementwise nonlinearities

    def __matmul__(self, other):
        other = as_tensor(other)
        out = Tensor(self.data @ other.data, (self, other))

        def _backward():
            self.grad += out.grad @ other.data.T
            other.grad += self.data.T @ out.grad
        out._backward = _backward
        return out

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

        def _backward():
            grad = out.grad
            if axis is not None and not keepdims:
                axes = (axis,) if isinstance(axis, int) else tuple(axis)
                for ax in sorted(axes):
                    grad = np.expand_dims(grad, ax)
            self.grad += np.ones_like(self.data) * grad
        out._backward = _backward
        return out

    def exp(self):
        out = Tensor(np.exp(self.data), (self,))

        def _backward():
            self.grad += out.grad * out.data
        out._backward = _backward
        return out

    def log(self):
        out = Tensor(np.log(self.data), (self,))

        def _backward():
            self.grad += out.grad / self.data
        out._backward = _backward
        return out

    def relu(self):
        out = Tensor(np.maximum(self.data, 0.0), (self,))

        def _backward():
            self.grad += out.grad * (self.data > 0.0)
        out._backward = _backward
        return out

The final piece is topological order. A depth-first search lists parents before users; reversing that list ensures every node has received all downstream adjoints before its backward closure runs.

Listing 10.3 Topological sort and backward pass

    def backward(self, gradient=None):
        if gradient is None:
            gradient = np.ones_like(self.data)
        topo, seen = [], set()

        def build(node):
            if id(node) in seen:
                return
            seen.add(id(node))
            for child in node._prev:
                build(child)
            topo.append(node)

        build(self)
        for node in topo:
            node.grad = np.zeros_like(node.data)
        self.grad = np.asarray(gradient, dtype=np.float64)
        for node in reversed(topo):
            node._backward()
        return self.grad

The tests build a tiny network using add, matrix multiply, relu, exp, log, multiplication, and sum, then compare the resulting gradients with central finite differences from scratch.gradcheck.check_gradient. The implementation is intentionally small: no mutation, no in-place operations, no convolutions, and no mixed precision. Those omissions keep the contract visible: each operation creates a value and a local backward closure.

10.5 Tapes and graphs

Frameworks differ mainly in when they record the graph. Eager systems record a tape while the Python program runs, which is easy to debug. Staged systems trace or compile a graph before execution, which enables larger compiler optimizations but makes Python control flow part of the tracing contract. Dynamic branches are fine only if the framework records or recompiles the path that actually ran. Both still rely on the same local VJPs and reverse topological sweep. Compilers can fuse operations, delete unused work, or rematerialize activations, but the mathematical object they preserve is the VJP of the original program.

In practice

Backpropagation made multilayer neural networks trainable [rumelhart1986], and modern autodiff systems are its programmable form [baydin2015automatic]. Large LLM training uses reverse mode because losses are scalar and parameter counts are huge. Memory, not algebra, is often the bottleneck: activation checkpointing recomputes selected forward values during the backward pass to reduce stored activations [chen2016training]. Distributed systems such as PyTorch FSDP combine autodiff with sharding so gradients, parameters, and optimizer states can fit across devices [zhao2023pytorch].

Key equations

\dot{\vv} = \mJ_f(\vu)\dot{\vu}

\bar{\vu} = \bar{\vv}\,\mJ_f(\vu)

z=x+y: \quad \bar{x} \mathrel{+}= \bar{z},\; \bar{y} \mathrel{+}= \bar{z}

z=xy: \quad \bar{x} \mathrel{+}= \bar{z}y,\; \bar{y} \mathrel{+}= \bar{z}x

\mY=\mA\mW: \quad \bar{\mA}=\bar{\mY}\mW^\T,\; \bar{\mW}=\mA^\T\bar{\mY}

10.6 Teach it

The one-sentence version: autodiff records how each value was computed, then applies local chain-rule rules backward from the loss to every parameter. Analogy: a receipt totals a bill forward, but a refund traces the total backward to each item that contributed. Board steps: draw a graph; write a local VJP for multiply; explain why shared nodes add adjoints; run nodes in reverse topological order. Misconceptions: autodiff is not finite differences; reverse mode is not symbolic simplification; gradients must be summed over broadcasted axes. Check for understanding: why does a scalar loss make reverse mode cheaper than running one forward-mode pass per parameter?

10.7 Exercises

Exercise 10.1 ★ Computational graph

Draw the graph for $L = \log(xy + x^2 + e^y)$ . Which node is shared, and why must its adjoint accumulate contributions?

Solution in Appendix D

Exercise 10.2 ★★ VJP derivation

Derive the VJP rules for $z = xy$ and for $\mY = \mA\mW$ . State the shapes of the matrix adjoints.

Solution in Appendix D

Exercise 10.3 ★★ Forward or reverse?

A function has many parameters and one scalar loss. Explain why reverse mode is preferred over forward mode. Give one case where forward mode would be attractive.

Solution in Appendix D

Exercise 10.4 ★★★ Implementation

Use the tiny Tensor to compute gradients for X @ W + b).relu().exp().log(.sum() and check them against finite differences.

Solution in Appendix D

References

[rumelhart1986] D. E. Rumelhart, G. E. Hinton, and R. J. Williams. Learning representations by back-propagating errors. Nature 323, 533–536, 1986.
[griewank2008] A. Griewank and A. Walther. Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, 2nd edition. SIAM, 2008.
[baydin2015automatic] A. G. Baydin et al. Automatic Differentiation in Machine Learning: a Survey. 2015. arXiv:1502.05767
[chen2016training] T. Chen et al. Training Deep Nets with Sublinear Memory Cost. 2016. arXiv:1604.06174
[zhao2023pytorch] Y. Zhao et al. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. 2023. arXiv:2304.11277