Capstone: a tiny autograd

A value and its history

The central object is a number with three extra fields. It has a gradient, a list of the values that produced it, and a function that knows how to push a gradient back to them. Every operation returns a new value and records the connection. Addition and multiplication are enough to build a surprising amount, and a tanh supplies the nonlinearity.

The forward pass is ordinary arithmetic with bookkeeping attached. When you multiply two values, the result stores its parents and the local derivative. Its parents are the two inputs, and the derivative with respect to one is the other. Nothing is computed yet. The graph is a record of what happened, waiting for a gradient to flow backward through it.

A scalar engine is enough to see the mechanism, because every value is one number. There is no shape to check and no broadcasting to reason about. The only things left are the chain rule and the order of the walk. Those are the right things to learn first: the bugs are conceptual, not dimensional.

Backward is a topological walk

To run backward, collect every node reachable from the output. Order them so that each node comes after its parents, then walk the reverse. Seed the output gradient with one. Each node calls its own backward function, which adds the incoming gradient times the local derivative to each parent. Addition accumulates, which is exactly what a shared node needs.

The order matters. If a node is visited before one of its consumers has sent a gradient, it forwards a partial value and the rest is lost. A depth-first post-order gives a valid order for a graph without cycles, and a forward pass always produces such a graph. A framework does the same walk, with tensors instead of scalars.

Cycles are impossible in a forward pass. Every operation creates a new value from existing ones, so time only moves forward. That is what makes the topological sort terminate. A recurrent network unrolled over time is still a finite acyclic graph. Acyclic means it has no loops. It has one copy of the cell per step. Backpropagation through time is therefore the same algorithm with a longer chain.

The same graph in PyTorch

import torch

w = torch.tensor([0.5, -0.5], requires_grad=True)
x = torch.tensor([1.0, 2.0])

loss = (w * x).sum().tanh()
loss.backward()

print(loss.item())
print(w.grad)   # d loss / d w

What a real engine adds

A production engine works on arrays rather than scalars. One node then represents a whole tensor, and its backward function is a matrix operation. The engine fuses kernels, runs on a GPU, reuses memory, and handles broadcasting by summing over the axes that were stretched. None of that changes the idea. It changes the constant factor by many orders of magnitude.

It also adds numerical care. A softmax is computed with the maximum subtracted. A log-sum-exp avoids overflow, and a division guards against a zero denominator. A scalar engine written for teaching can ignore those details, because the numbers are small. A real one cannot. That care is most of the difference between a demo and a library.

Memory is the other difference. A production engine frees intermediate buffers as soon as their gradient has been consumed. It can also recompute an activation instead of storing it when memory is tight. Checkpointing trades compute for memory at the granularity of a block. The scalar engine keeps everything, which is fine for four examples and impossible for a real batch.

From engine to network

A network is a function that composes values. Its parameters are leaves that no operation produced. Training is the same loop as before: forward, loss, backward, update, clear. The only new requirement is that the parameters are reachable from the loss through the graph. They are, if every layer used them.

This is the whole stack. Batching, optimisers and regularisation are additions around this loop, not replacements for it. Once the scalar version trains an XOR, a framework is a faster implementation of the same idea, not a black box.

The natural next step is to replace the scalar with a small array class and let the operations work elementwise. Broadcasting and matrix multiplication are the two additions that make it useful. Both are more about shape bookkeeping than about calculus. The backward rules stay the same. Only the data they move grows.