Autograd from scratch

Every deep learning framework rests on one call: loss.backward(). I wanted to know exactly what happens inside it, so I wrote the engine myself, without libraries. The code is in ML-from-scratch: v1 works on scalars, v2 on NumPy arrays.

A program is a graph

Any computation built from elementary operations is a directed acyclic graph. Take a squared error for a one-variable linear model:

y_hat = w * x + b
loss = (y_hat - y) ** 2

Each intermediate value is a node, and each node remembers which nodes it was computed from. With $w = 0.5$, $x = 2$, $b = 1$, $y = 3$:

The one rule

Write $\bar{v} = \partial L / \partial v$ for the gradient of the loss with respect to a node, its adjoint. The multivariable chain rule says a node's adjoint is the sum, over every node that consumed it, of that consumer's adjoint times the local derivative:

That is all of reverse-mode differentiation. Each operation only needs to know its own local derivatives. The sum is why gradients are accumulated with +=: a value used twice receives gradient from both uses.

Running the rule backward through the table:

One backward pass produces the gradient with respect to every input. Its cost is a small constant multiple of the forward pass, no matter how many parameters there are. This is why training uses reverse mode rather than forward mode, which would need one pass per parameter.

The node

class Value:
    def __init__(self, value, parents=()):
        self.value = value
        self.parents = parents
        self.grad = 0.0
        self._backward = lambda: None
 
    def __mul__(self, other):
        out = Value(self.value * other.value, (self, other))
 
        def _backward():
            self.grad += other.value * out.grad
            other.grad += self.value * out.grad
 
        out._backward = _backward
        return out

The other operations differ only in their local derivatives:

Order matters

A node's adjoint is complete only after every consumer has added to it. So the backward pass visits nodes in reverse topological order, consumers before the values they consumed. A depth-first search gives the order:

def backward(self):
    order, seen = [], set()
 
    def visit(v):
        if v not in seen:
            seen.add(v)
            for p in v.parents:
                visit(p)
            order.append(v)
 
    visit(self)
    self.grad = 1.0
    for v in reversed(order):
        v._backward()

Training

With gradients available, gradient descent is one line per parameter, $\theta \leftarrow \theta - \eta\, \bar\theta$, followed by resetting the gradient so the next step does not inherit this one.

The v1 data is $y = x - 1$ for $x = 0, \dots, 99{,}999$, with both columns divided by $10^5$. The model should learn and , and it converges to both.

Matrices

v2 keeps the same graph and the same rule; only the local derivatives change. Two of them are worth writing down.

Matrix multiplication. For $Z = AB$ with upstream gradient $\bar{Z}$:

The shapes force this form: $\bar{A}$ must have $A$'s shape, and $\bar{Z} B^\top$ is the only product of the available matrices that does.

Broadcasting. Adding a bias $b \in \mathbb{R}^{n}$ to every row of $X \in \mathbb{R}^{m \times n}$ uses $b$ $$ times. By the sum in the rule above, its gradient is the upstream gradient summed over the broadcast axis:

In general, a gradient is summed over every axis that broadcasting added or stretched, until it matches the shape of the original value. Forgetting this is the most common bug in hand-written backward passes, and v2 spends most of its gradient code on getting it right.

What I took from it

Backpropagation is not a separate algorithm layered on calculus. It is the chain rule, evaluated in an order that lets each local derivative be used once. Frameworks add tensors, kernels, memory planning and fused operations, but torch.Tensor.backward() still runs the loop above.

References

1. Baydin et al. Automatic Differentiation in Machine Learning: a Survey. JMLR 2018. arXiv:1502.05767
2. Karpathy. micrograd.
3. Olah. Calculus on Computational Graphs: Backpropagation. 2015.