-
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 dLossdp\frac{dLoss}{dp}dpdLoss. 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;So the entire problem reduces to one question: how do we get dLossdp\frac{dLoss}{dp}dpdLoss 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, a+ba + ba+b: nudge aaa up by one, the result goes up by one. So ∂(a+b)∂a=1\frac{\partial(a+b)}{\partial a} = 1∂a∂(a+b)=1.
- for multiplication, a⋅ba \cdot ba⋅b: nudge aaa, the result changes by bbb. So ∂(a⋅b)∂a=b\frac{\partial(a \cdot b)}{\partial a} = b∂a∂(a⋅b)=b — 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:dzdx=dzdy⋅dydx\frac{dz}{dx} = \frac{dz}{dy} \cdot \frac{dy}{dx}dxdz=dydz⋅dxdy
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:dzdx=(dzdx)path 1+(dzdx)path 2+⋯\frac{dz}{dx} = \left(\frac{dz}{dx}\right)_{\text{path 1}} + \left(\frac{dz}{dx}\right)_{\text{path 2}} + \cdotsdxdz=(dxdz)path 1+(dxdz)path 2+⋯
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); Value *L = value_mul(e, a);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 e=a⋅be = a \cdot be=a⋅b, then L=e⋅aL = e \cdot aL=e⋅a. Substituting, L=a⋅b⋅a=a2⋅bL = a \cdot b \cdot a = a^2 \cdot bL=a⋅b⋅a=a2⋅b. With a=2a = 2a=2 and b=3b = 3b=3, that's L=12L = 12L=12.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 dLda\frac{dL}{da}dadL and dLdb\frac{dL}{db}dbdL. Starting from the output and chaining backwards:
- L=e⋅aL = e \cdot aL=e⋅a, so the local derivatives are ∂(e⋅a)∂e=a=2\frac{\partial (e \cdot a)}{\partial e} = a = 2∂e∂(e⋅a)=a=2 and ∂(e⋅a)∂a=e=6\frac{\partial (e \cdot a)}{\partial a} = e = 6∂a∂(e⋅a)=e=6.
- e=a⋅be = a \cdot be=a⋅b, so ∂e∂a=b=3\frac{\partial e}{\partial a} = b = 3∂a∂e=b=3 and ∂e∂b=a=2\frac{\partial e}{\partial b} = a = 2∂b∂e=a=2.
Now assemble them with the chain rule.
bis easy — it only reachesLthroughe:dLdb=dLde⋅∂e∂b=a⋅a=2⋅2=4\frac{dL}{db} = \frac{dL}{de} \cdot \frac{\partial e}{\partial b} = a \cdot a = 2 \cdot 2 = 4dbdL=dedL⋅∂b∂e=a⋅a=2⋅2=4
ais the interesting one — it reachesLthrough two paths, so we add them:dLda=(dLde⋅∂e∂a)+(∂L∂a)directly=a⋅b+e=2⋅3+6=12\begin{aligned} \frac{dL}{da} &= \left(\frac{dL}{de} \cdot \frac{\partial e}{\partial a}\right) + \left(\frac{\partial L}{\partial a}\right)_{\text{directly}} \\ &= a \cdot b + e \\ &= 2 \cdot 3 + 6 = 12 \end{aligned}dadL=(dedL⋅∂a∂e)+(∂a∂L)directly=a⋅b+e=2⋅3+6=12
Six from the path through eee, 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 aaa, and push its influence forward through the graph: compute deda\frac{de}{da}dade, then dLda\frac{dL}{da}dadL. One sweep gives you the derivative of everything with respect to aaa. But you only learned about aaa. To also learn about bbb, you'd do another whole sweep. One sweep per input.
Backward. Start at the output LLL, seed it with dLdL=1\frac{dL}{dL} = 1dLdL=1, and push influence backwards toward the inputs. One sweep fills in dLde\frac{dL}{de}dedL, dLda\frac{dL}{da}dadL, and dLdb\frac{dL}{db}dbdL — 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 ML world. Everything microcrad does after building the graph is a single backward sweep.
Note that, before a node's gradient can feed the nodes behind it, that gradient has to be finished. For instance, let's look back at dLda\frac{dL}{da}dadL: it wasn't done until both the path through e and the direct path had integrated their contribution. If the nodes are visited in the wrong order, you'll propagate a half-summed gradient. To avoid this, we sort the graph so each node comes after the nodes it depends on in a list, then walk that list back to front — that's reverse topological order, and it's the first step of the backward pass according to the backpropagation algorithm.
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; double grad; struct Value **prev; int32_t op_code; } 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:
- The result of the operation
v1->data * v2->datais computed and stored in a fresh node. - The operands are recorded in
prev. - The resulting node is tagged with the operation that made it (
MUL_OP_CODEin this case), 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:a = 2b = 3e = a * b = 6L = e * a = 12
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 dLda\frac{dL}{da}dadL 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. Naturally, the loss becomes the output node of the graph.
- 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; 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; } }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 (a+ba + ba+b): local derivative is 111, so the gradient passes straight through, unchanged.
- mul (a⋅ba \cdot ba⋅b): 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); Value *L = value_mul(e, a); value_backward(L); 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.000000dLda=12\frac{dL}{da} = 12dadL=12 and dLdb=4\frac{dL}{db} = 4dbdL=4 — exactly what we got with pencil and paper. Look at the 121212: it's the 666 from the path through eee plus the 666 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, so 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 w⋅x+bw \cdot x + bw⋅x+b; stack neurons into layers and layers into a neural network, 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++) vector_get(params, i)->grad = 0.0; value_backward(loss); for (size_t i = 0; i < params->size; i++) { Value *p = vector_get(params, i); p->data -= learning_rate * p->grad; }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 y=max(0,x1+x2)y = \max(0, x_1 + x_2)y=max(0,x1+x2). 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 L=a2⋅bL = a^2 \cdot bL=a2⋅b.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 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.
oraziorillo.com
Technology
What loss.backward() actually does
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...