-
If you've trained a neural network before, chances are you've typed
loss.backward()without being able to say what exactly happens under the hood. By the end of this post you'll understand the core mechanism behind engines like PyTorch well enough that you could write it yourself.To keep things hands-on, we'll often refer to microcrad — a scalar-valued automatic differentiation engine I recently wrote, inspired by Andrej Karpathy's micrograd.
You don't need to know C to follow along. If you can read simple code and have some basic knowledge of neural networks and calculus, you'll be fine.
What are we even computing?
Training a neural network means minimizing a loss: a single number that measures how wrong the network currently is. You minimize it with gradient descent: nudge every parameter a little bit in the direction that makes the loss go down, then repeat. To do that, for each parameter
pyou need to know:if I wiggle
pa tiny bit, how much does the loss change, and in which direction?That quantity is the derivative . Compute it for every parameter, take a small step against it, and you've done one step of learning:
p->data -= learning_rate * p->grad; /* one gradient descent step */So the entire problem reduces to one question: how do we get for every
pat once? A real network has thousands to trillions of them, and there's exactly one loss. Hold on to that shape — many inputs, one output — because it's the reason everything later works the way it does.One operation at a time: local derivatives and the chain rule
You don't need to know the derivative of your whole monstrous network. You only need the derivative of each individual operation, in isolation. For the two operations we'll keep using in this post:
- for addition, : nudge up by one, the result goes up by one. So .
- for multiplication, : nudge , the result changes by . So — the other operand.
Every other operator — subtraction, hyperbolic tangent, whatever — works the same way. For simplicity, I'll only show addition and multiplication and let the rest be variations on these two.
These are local derivatives: how one operation's output moves when you nudge one of its direct inputs, holding the rest fixed.
The interesting part is gluing them together, and that's the chain rule. If
zdepends ony, andydepends onx, then:Derivatives compose by multiplying along the path. To learn how
xaffects a farawayz, you walk the path fromzback toxand multiply the local derivatives you pass through. That's the engine of backpropagation.There's one more clause. When
xreacheszthrough more than one path, the contributions add up:Our running example
Here's the example we'll carry for the rest of the post:
Value *a = value_create_leaf(2.0); Value *b = value_create_leaf(3.0); Value *e = value_mul(a, b); /* e = a * b = 6 */ Value *L = value_mul(e, a); /* L = e * a = 12 */Valueis microcrad's one and only fundamental type — it wraps a single double precision number. Ignore the exact function names for a second and just read the math: we compute , then . Substituting, . With and , that's .Notice that
ais used twice — once to makee, and once directly inL. That's the multiple-paths case from above, hiding in four lines of code. Let's differentiateLby hand.We want and . Starting from the output and chaining backwards:
- , so the local derivatives are and .
- , so and .
Now assemble them with the chain rule.
bis easy — it only reachesLthroughe:ais the interesting one — it reachesLthrough two paths, so we add them:Six from the path through , six from the direct path, twelve total.
Why we go backwards
We now have all the pieces to compute a derivative. But how we compute them matters enormously.
There are two directions you could apply the chain rule.
Forward. Pick one input, say , and push its influence forward through the graph: compute , then . One sweep gives you the derivative of everything with respect to . But you only learned about . To also learn about , you'd do another whole sweep. One sweep per input.
Backward. Start at the output , seed it with , and push influence backwards toward the inputs. One sweep fills in , , and — the derivative of the output with respect to everything. One sweep, all inputs.
Now remember the shape of our problem: many parameters, one loss. Forward mode costs one sweep per parameter — catastrophic when you have a million of them. Backward mode costs one sweep, and hands you the gradient for every parameter at once. That asymmetry is the entire reason neural networks are trainable at all.
This is reverse-mode automatic differentiation, and "backpropagation" is just its name in the neural-network world. Everything microcrad does after building the graph is a single backward sweep.
One catch makes it work. Before a node's gradient can feed the nodes behind it, that gradient has to be finished. Look back at : it wasn't done until both the -path and the direct path had chipped in. Visit the nodes in the wrong order and you'll read a half-summed gradient and propagate garbage. The fix is to visit every node before any of its operands. Sort the graph so each node comes after the nodes it depends on, then walk that list back to front — that's reverse topological order, and it's why the backward pass starts by sorting the graph.
The graph builds itself
So the backward pass needs a graph — a record of which operation produced which value from which operands. But you never build that graph explicitly; you just do the forward computation, and the graph falls out as a side effect.
The secret is that a
Valueis not just a number. It's a number that remembers where it came from:typedef struct Value { double data; /* the scalar this node holds */ double grad; /* dLoss/dThisValue, filled in by the backward pass */ struct Value **prev; /* the operands (previous nodes in the graph) */ int32_t op_code; /* which operation produced this Value */ /* memory-management fields are omitted */ } Value;A
Valueproduced by an operation points back at the operands it was computed from throughprev, and tags itself with the operation viaop_code. Follow thoseprevpointers from any node and you're walking the computation graph backwards.But how does a node get wired up? Look at what a single operation does. Here's multiplication, with the error handling stripped out:
Value *value_mul(Value *v1, Value *v2) { Value **prev = malloc(2 * sizeof(Value *)); prev[0] = v1; value_retain(v1); prev[1] = v2; value_retain(v2); Value *result = value_create(v1->data * v2->data, 2, prev); result->op_code = MUL_OP_CODE; return result; }Three things happen, and the same three happen for every operation in the engine:
- It computes the result of the operation
v1->data * v2->dataand stores it in a fresh node. - It records which operands produced that number, by stashing them in
prev. - It tags the node with the operation that made it (
MUL_OP_CODE), so the backward pass will later know which derivative rule to apply.
(The two
value_retaincalls are bookkeeping for C's manual memory management. Ignore them for now, I'll come back to them.)So when you wrote those four lines of our running example, you weren't just computing
12. Every operation quietly left behind a node pointing at its inputs, and by the time you hadLin hand, you were also holding the root of a graph recording the entire history of howLcame to be:Look carefully at
a: it's a single node with two arrows coming out of it — one intoe, one straight intoL. That fork is exactly the "two paths" we differentiated by hand, and it's why will have two contributions to add. That history is everything the backward pass needs.The backward pass, in code
value_backwarddoes exactly the two steps we reasoned our way to:- Build a topological ordering of the graph, so every node comes after the nodes it depends on.
- Seed the output's gradient to
1and walk that list in reverse, applying one local-derivative rule per node.
Step two is a single
switchover the operation codes. The full version has a few more cases, but the important part is the shape, not the catalog:v->grad = 1; /* stops before index 0: topo->data[0] is always a leaf — nothing to propagate */ for (size_t i = topo->size - 1; i != 0; i--) { Value *node = topo->data[i]; switch (node->op_code) { case ADD_OP_CODE: for (uint32_t j = 0; j < node->n_prevs; j++) node->prev[j]->grad += node->grad; break; case MUL_OP_CODE: for (uint32_t j = 0; j < node->n_prevs; j++) node->prev[j]->grad += node->grad * node->prev[1-j]->data; break; /* ... same pattern for the other ops ... */ } }Every case does the same high-level thing: take the gradient already sitting on the node (
node->grad), multiply by the operation's local rule, and accumulate the result into the operands:- add (): local derivative is , so the gradient passes straight through, unchanged.
- mul (): local derivative for one operand is the other operand.
Watching it work
Let's push the example through the real engine and ask it for the gradients we computed by hand:
Value *a = value_create_leaf(2.0); Value *b = value_create_leaf(3.0); Value *e = value_mul(a, b); /* e = a * b = 6 */ Value *L = value_mul(e, a); /* L = e * a = 12 */ value_backward(L); /* fills in every grad field */ printf("dL/de = %f\n", e->grad); printf("dL/da = %f\n", a->grad); printf("dL/db = %f\n", b->grad);Compiled and run against microcrad, this prints:
dL/de = 2.000000 dL/da = 12.000000 dL/db = 4.000000and — exactly what we got with pencil and paper. Look at the : it's the from the path through plus the from the direct path, summed together.
What Python hides
I promised I'd come back to those
value_retaincalls. In Python, you build the graph, call backward, and let the garbage collector clean up. In C there's no collector, and a computation graph is a tangle of shared pointers — the same weight can be an operand of thousands of nodes.Microcrad handles this with reference counting: every
Valuecounts how many things point at it, each operation retains its operands, and when a count hits zero the node frees itself and releases its operands in turn. Release the root of a graph and the whole thing cascades away, freeing exactly the intermediates nothing else holds while the weights survive. It's a satisfying property, but it has little to do with backpropagation, so you can happily file it as "the thing that keeps the graph from leaking". The details are in the README of the repo.That's it
That's the whole idea. To turn it into learning, you wrap it in a loop. A
Value-based neuron computes an activation on top of ; stack neurons into layers and layers into an MLP, and because every weight is aValue, a forward pass automatically builds a graph you can backpropagate through. Then every training step is the same four moves:for (size_t i = 0; i < params->size; i++) /* 1. zero old gradients */ vector_get(params, i)->grad = 0.0; value_backward(loss); /* 2. backprop: fill every grad */ for (size_t i = 0; i < params->size; i++) { /* 3. step against the gradient */ Value *p = vector_get(params, i); p->data -= learning_rate * p->grad; } /* 4. release the graph and repeat */That's the entire loop that trains everything from a toy regression to GPT. The tensors get bigger and the hardware gets fancier, but the four steps don't change. Microcrad ships a tiny example that trains a small network to learn . Run
make example_toy_regressionand watch the loss fall.Where to go from here
If you've read this far, you know (almost) exactly what
loss.backward()does: it builds a graph as a side effect of the forward pass, sorts it so every node comes after its dependencies, then walks it in reverse, scattering each node's gradient onto its operands with a+=that quietly implements the multivariable chain rule.The best way to make it stick is to stop reading and go implement one more operator yourself, e.g. sigmoid. Fork the repo, add the forward op, add one
caseto theswitch, and check the gradient with a pencil as we did with .To be clear about the gap: microcrad is scalars only, a handful of ops, CPU only. What PyTorch adds on top — tensors, GPU kernels, hundreds of operations — is a mountain of engineering, but not a different idea. The math of backpropagation really is as simple as it looks. The only thing standing between "I can use PyTorch" and "I can implement (one small fundamental piece of) PyTorch" is sitting down and writing it once. microcrad was my attempt to get one step closer to that.