Reverse-Mode Differentiation on the Computational Graph

Training needs the derivative of one scalar, the loss, with respect to millions of parameters. Written out, that derivative is the product of as many Jacobians as the network has operations, and forming any of those Jacobians explicitly is out of the question: a dense layer mapping a thousand units to a thousand units has a $10^6$-entry Jacobian, and a batch of a hundred samples implies a hundred of them. Reverse-mode differentiation never forms one. It evaluates the chain rule as a sequence of matrix-vector products, in the reverse order of the forward computation, using only quantities that already appear in the forward pass.

Building the tape during the forward pass

The forward pass is executed as a sequence of primitive operations, each of which takes one or more values and produces one value. Each primitive has a known local derivative rule. During the forward pass the system records, for every produced value, which operands it came from and a reference to the numbers needed to evaluate the local derivative rule at that point. That record is what is loosely called the tape or the graph.

The local derivative is not usually evaluated during the forward pass; it is evaluated later, from saved operands or saved outputs. This is a deliberate trade. For an elementwise activation the rule needs the pre-activation tensor, so the pre-activation is retained; for a matrix product the rule needs both factors, so both are retained. Recomputing them instead of storing them is possible and is what the memory-for-compute trade in activation checkpointing does. Because the recorded graph is a directed acyclic graph whose edges point from operands to results, the reverse of the order in which values were created is a valid topological order, which is why an implementation can walk a simple list rather than sort a graph.

The reverse sweep and the sum over consumers

The sweep visits nodes in reverse creation order. Each node holds a cotangent, the derivative of the final scalar with respect to that node’s value. The output node is seeded with a cotangent of 1, and each node distributes its cotangent to its operands, scaled by the local derivative along that edge. When a value has more than one consumer, the contributions add:

$$\frac{\partial L}{\partial x} = \sum_{k} \frac{\partial L}{\partial y_k} \frac{\partial y_k}{\partial x}$$

where $y_1, \ldots, y_k$ are the consumers of $x$. This accumulation is the part that hand-written versions most often get wrong, because the natural data structure for storing cotangents is a dictionary keyed by value, and assigning into that dictionary overwrites the previous contribution instead of adding to it. The symptom is a gradient that is correct for the last consumer encountered in the sweep and missing the rest, which shows up as a model that trains slightly worse than it should rather than as an exception. Weight tying produces exactly this pattern: one matrix used as both an input embedding and an output projection has its gradient contributions summed over both uses, as noted in dense layers.

class Value:
    def __init__(self, data, parents=(), local=()):
        self.data = data
        self.parents = parents          # operands this value was computed from
        self.local = local              # d(self)/d(parent), one entry per parent
        self.grad = 0.0

    def backward(self):
        order, seen = [], set()
        def visit(v):
            if id(v) in seen:
                return
            seen.add(id(v))
            for p in v.parents:
                visit(p)                     # operands precede the result
            order.append(v)
        visit(self)
        self.grad = 1.0                      # seed the sweep at the scalar output
        for v in reversed(order):
            for p, d in zip(v.parents, v.local):
                p.grad += v.grad * d         # fan-out adds; it must not assign
        return self

def mul(a, b):
    return Value(a.data * b.data, (a, b), (b.data, a.data))

def add(a, b):
    return Value(a.data + b.data, (a, b), (1.0, 1.0))

x, y = Value(3.0), Value(4.0)
loss = add(mul(x, y), x)                     # loss = x*y + x  =>  dloss/dx = y + 1
loss.backward()
assert x.grad == 5.0 and y.grad == 3.0

Why the cost is independent of the parameter count

Every primitive’s vector-Jacobian product costs the same order of work as its forward evaluation, so one reverse sweep over the graph costs a constant multiple of the forward pass, and that multiple does not depend on how many parameters exist. A dense layer’s forward pass costs $N n_{\text{in}} n_{\text{out}}$ multiply-accumulates and its contribution to the reverse sweep costs about twice that, to produce gradient contributions for both the weights and the input. The same constant also does not depend on the number of inputs, only on the number of outputs of the fused function, which is why the sweep is efficient for a scalar loss and would degrade to $K$ separate seeds for a $K$-dimensional output vector. That direction-of-differentiation choice is treated in Jacobian-vector and vector-Jacobian products. Forward-mode differentiation, by contrast, computes one directional derivative per seed direction, so a full gradient with respect to $p$ parameters would need $p$ sweeps.

The cost of the sweep is paid in memory rather than arithmetic. Every intermediate whose local derivative rule needs it must remain addressable until the reverse sweep reaches it, so peak memory grows with the number of recorded tensors and with the size of each, which for a batch of activations is proportional to the batch size. A layer producing 1024 activations for 64 examples holds $64 \times 1024 = 65536$ values, which is 256 KiB in 32-bit floats, and a stack of a hundred such layers retains roughly 25 MiB of intermediate state for one tensor per layer, before any larger intermediate such as a full batch of feature maps is counted. Recomputing the forward pass for a segment of the network during the sweep instead of retaining its intermediates replaces that memory with additional forward work, at the cost of evaluating part of the network twice; the choice between the two is a memory budget decision, not a correctness one, since the numbers the local rules read are the same either way. Freeing a tensor as soon as the sweep consumes it is safe and standard; freeing it before the sweep reaches it is not.

Where the sweep silently produces wrong gradients

Three classes of defect recur. The first is mutation: if a saved tensor is modified after the forward pass, the local derivative rule is evaluated against values that no longer correspond to the forward computation, and the resulting gradient is wrong without any error being raised. In-place operations on tensors whose result is needed for the forward pass are safe only when the system can prove that nothing else needs the overwritten value.

The second is nondifferentiable primitives. Operations such as absolute value, maximum, and integer indexing are continuous but not differentiable at some points, and the rules applied there pick one of the possible one-sided derivatives. A maximum over a set of values attributes the gradient to whichever argument was largest, so ties split arbitrarily, and a value that never wins a comparison receives exactly zero gradient, which can look like a dead unit.

The third is control flow that depends on data. The recorded graph contains only the branch that was executed, so the derivative computed by the sweep is the derivative of that branch. Differentiating the sweep itself to obtain second derivatives through a branch that changes with the input gives the derivative of a function that no longer describes the computation. Repeated calls to the sweep also accumulate into the existing gradient buffers rather than replacing them, which is why clearing gradients before each step is part of the interface, and why a second sweep over a graph whose saved tensors have already been consumed fails instead of returning a number.

Accumulation, reduction, and floating-point order

Because the sweep sums contributions, its numerical result depends on the order of operations within a primitive and on how reductions are laid out in memory, so two implementations that are mathematically identical can differ in the last few bits. This matters for gradient checking, where a relative difference of $10^{-7}$ in single precision is unremarkable and a difference of $10^{-3}$ is not. The practical consequence is that a derivative computed in 32-bit arithmetic should be compared with finite differences at a tolerance derived from that precision, not from the precision of the difference formula, and the choice of step size in such a check follows from the same constraint.