The Chain Rule and Computational Graphs: How Gradients Really Flow

A modern neural network can have billions of parameters, and every single one of them gets nudged in the right direction after each batch of data. How does the computer know which way to nudge each one? The answer is one idea from calculus, the chain rule, organised into a structure called a computational graph. Today we build both from scratch.

• Every network is a chain
• Gradients flow backward

In the last lesson we learned to split data without leakage and wrap everything in pipelines. Now we go back to the mathematics. This lesson is the foundation under lesson 26, where you saw backpropagation, and lesson 49, where we meet PyTorch autograd. We will not re-teach a full neural network. Instead we make sure the engine underneath is crystal clear.

• Lesson 35: leakage-safe splits
• Lesson 36: chain rule

Here is the plan. First we derive the chain rule from the idea that rates multiply, and check it with numbers. Then we extend it to many variables, where contributions add up over paths. Next we draw a computational graph for a squared error loss and compute every gradient by hand. Finally we handle fan-out, view it through Jacobians, and write a working autograd class in about forty lines of Python.

• Rates multiply
• Graphs and reverse pass
• Code your own autograd

Imagine three gears in a row. When gear x turns once, gear u turns three times. When gear u turns once, gear y turns fourteen times. So how many times does gear y turn when x turns once? Three times fourteen, which is forty two. Rates of change along a chain simply multiply, and that is the entire chain rule in one picture.

• x drives u
• u drives y

Now the formal version. A composite function feeds the output of one function into another. We call the inner function g, which turns x into u, and the outer function f, which turns u into y. Written together, y equals f of g of x. In our example, u equals three x plus one, and y equals u squared. At x equal to two, u is seven and y is forty nine.

• y = f(u)
• u = g(x)
• y = f(g(x))

Let us derive it. Nudge x by a tiny amount, delta x. That moves u by delta u, which in turn moves y by delta y. Now write the ratio delta y over delta x as delta y over delta u, times delta u over delta x. The delta u cancels, which is just fractions. Let delta x shrink toward zero, and each ratio becomes a derivative, giving dy by dx equals dy by du times du by dx.

• Nudge x by delta x
• u moves by delta u
• y moves by delta y

A careful mathematician notices one gap: if delta u happens to be exactly zero, we divided by zero. Rigorous proofs patch this with a slightly different argument, and the result is the same, as long as both functions are differentiable. For machine learning, the practical takeaway is simple. As long as each step is smooth, local slopes multiply along the chain, and we will use that without worry.

• Needs differentiable functions
• delta u must not be zero

This gives us a recipe. Break the function into simple pieces. Differentiate each piece on its own, treating its input as a plain variable. Then multiply. The outer piece, u squared, has derivative two u. The inner piece, three x plus one, has derivative three. Each of these is a local derivative, and the chain rule says the global derivative is their product.

• Differentiate each piece alone
• Multiply the results

Let us do the full calculation. At x equal to two, u is three times two plus one, which is seven, and y is seven squared, forty nine. Now the local derivatives. The slope of y with respect to u is two times seven, which is fourteen. The slope of u with respect to x is three. Multiply them: fourteen times three is forty two. So near x equal to two, y changes about forty two times as fast as x.

• x = 2, u = 7, y = 49
• dy/du = 14, du/dx = 3
• dy/dx = 42

Can we trust that? Let us check with a completely different route. Expand three x plus one, squared, to get nine x squared plus six x plus one. Differentiating term by term gives eighteen x plus six. At x equal to two, that is thirty six plus six, which is forty two. The two methods agree. The chain rule did not need the expansion, which matters enormously when functions are too tangled to expand.

• y = 9x² + 6x + 1
• dy/dx = 18x + 6
• At x = 2: 42

There is a second check that works for any code, called a finite-difference check. Evaluate the function at x equals two point zero zero one. Then u is seven point zero zero three, and y is forty nine point zero four two zero zero nine. Subtract forty nine, divide by zero point zero zero one, and you get forty two point zero zero nine. Very close to forty two. This is how engineers test gradient code.

• Nudge x by h = 0.001
• Compute the slope numerically

A better version is the central difference. Evaluate at x plus h and at x minus h, subtract, and divide by two h. Here we get forty nine point zero four two zero zero nine minus forty eight point nine five eight zero zero nine, over zero point zero zero two, which is exactly forty two. The classic mistake is stopping after the outer derivative, writing two times three x plus one, which is fourteen, and forgetting to multiply by the inner derivative, three.

• Use f(x+h) minus f(x−h)
• Forgetting the inner derivative

Real models have many variables. Suppose z depends on x and y, and both x and y depend on a single quantity t. Changing t moves x and also moves y, and both of those movements affect z. There are two different routes from t to z. A nudge in t travels down both routes at once, and the total effect on z is the combination of the two.

• z depends on x and y
• x and y depend on t

The rule is beautifully simple. Multiply derivatives along each path, then add the paths together. The first path goes through x, giving the partial of z with respect to x, times dx by dt. The second goes through y, giving the partial of z with respect to y, times dy by dt. Add them, and you have dz by dt. Multiply along, add across.

• Multiply along each path
• Add across paths

A concrete case. Let z equal x times y, with x equal to t squared and y equal to three t plus one. At t equal to two, x is four and y is seven. Path one: the partial of z with respect to x is y, which is seven, times dx by dt, which is two t, so four. That gives twenty eight. Path two: the partial with respect to y is x, which is four, times three, giving twelve. The total is forty.

• z = x·y, x = t², y = 3t + 1
• t = 2 gives x = 4, y = 7
• dz/dt = 28 + 12 = 40

Check it directly. Substituting gives z equal to t squared times three t plus one, which expands to three t cubed plus t squared. The derivative is nine t squared plus two t. At t equal to two, that is thirty six plus four, which is forty. It matches the sum over paths. Notice that if you dropped either path, you would get the wrong answer, either twenty eight or twelve. Forgetting a path is the multivariable version of forgetting the inner derivative.

• z = 3t³ + t²
• dz/dt = 9t² + 2t
• At t = 2: 40

Now we make this systematic. Any calculation a computer performs can be broken into tiny operations, each taking a few numbers and producing one. Draw each intermediate value as a node, and draw arrows showing which values feed which operations. The result is a computational graph. Complex formulas become a network of very simple steps, and each step has a derivative we already know.

• Nodes are values
• Edges are operations

Our running example is the squared error for a single data point from linear regression, as in lesson 14. The model predicts w times x plus b. The target is y. The error is the prediction minus the target, and we square it. That single number, f, is our loss. Training means finding how f changes when we nudge w and b, so those are the gradients we want.

• f = (w·x + b − y)²
• Squared error for one example

Break the formula into four elementary steps. First, a equals w times x. Second, c equals a plus b. Third, r equals c minus y, which is the residual. Fourth, f equals r squared. Every node holds one number, and every operation is something whose derivative we know by heart: a product, a sum, a difference, and a square.

• a = w·x
• c = a + b
• r = c − y
• f = r²

Let us run it with numbers. Say w equals two, x equals three, b equals one, and the target y equals ten. Forward pass: a is two times three, six. Then c is six plus one, seven. The residual r is seven minus ten, negative three. Finally f is negative three squared, nine. We store every one of these values, because the backward pass will need them.

• w = 2, x = 3, b = 1, y = 10
• a = 6, c = 7, r = −3, f = 9

Each operation only needs to know its own local derivatives. The product node has slope x with respect to w, which is three, and slope w with respect to x, which is two. The addition node has slope one for both inputs. The subtraction node has slope one for c and negative one for y. The square node has slope two r, which is negative six. No node ever needs to see the whole graph.

• Each node knows only its neighbors
• Compute them at the stored values

Now run the graph backward. At the output, the derivative of f with respect to itself is one. That is our starting seed. At every node, the gradient flowing to an input equals the upstream gradient, the one arriving from the output side, times the local derivative of that node. It is the chain rule, applied one edge at a time, moving from the loss back toward the inputs.

• Gradient = upstream × local
• Start with df/df = 1

Step by step. The seed at f is one. The square node multiplies by negative six, so the gradient at r is negative six. The subtraction node passes that on to c unchanged, negative six, and flips the sign for y, giving positive six. The addition node passes negative six to both a and b. So far we have the gradient with respect to every intermediate value.

• df/dr = −6
• df/dc = −6, df/dy = 6
• df/da = −6, df/db = −6

The last node is the product. The upstream gradient is negative six. The gradient for w is negative six times x, which is three, so negative eighteen. The gradient for x is negative six times w, which is two, so negative twelve. The bias gradient was already negative six. Notice the sign: negative gradients mean increasing w or b would lower the loss.

• df/dw = −6 × 3 = −18
• df/dx = −6 × 2 = −12

Check against the direct formula. The derivative of f with respect to w is two times the residual times x, which is two times negative three times three, negative eighteen. It matches. Now one gradient descent step from lesson 11 with learning rate zero point zero one: w becomes two minus zero point zero one times negative eighteen, which is two point one eight. The new residual is about negative two point four six, and the loss drops from nine to about six.

• Formula: 2r·x = −18
• w ← w − 0.01·(−18) = 2.18

What if one value is used in several places? Take f equals x times x plus three x, at x equal to two. The variable x fans out into the product node, twice, and into the times three node. On the way back, each use sends its own gradient to x, and all of them must be added. This is exactly the sum over paths rule from earlier, now built into the algorithm.

• A value used twice
• Gradients from each use add

The numbers. Both branches receive upstream gradient one from the sum node. The product x times x sends one times x, which is two, to each of its two input slots, so x gets two plus two, four. The times three node sends three. The total is seven. Compare with the derivative of x squared plus three x, which is two x plus three, equal to seven. A very common bug is overwriting a gradient instead of adding to it.

• Square node: 2 + 2 = 4
• Times-three node: 3
• Total = 7

Back in lesson 7 we met the Jacobian, the matrix of all partial derivatives of a vector function. It is the best linear approximation of that function near a point. Take g that maps x one, x two to u one equal to x one plus x two, and u two equal to x one times x two. At the point two, three, the Jacobian rows are one, one and three, two.

• Vector in, vector out
• Jacobian: matrix of partials
• Composition: multiply matrices

The vector chain rule says the Jacobian of a composition is the product of the Jacobians, in order. Let f of u equal u one squared plus u two. At u equal to five, six, its gradient row is ten, one. Multiply that row by the Jacobian of g: ten times one plus one times three gives thirteen, and ten times one plus one times two gives twelve. Direct differentiation of x one plus x two, squared, plus x one x two confirms thirteen and twelve.

• J of f∘g = J_f · J_g
• Shapes: k×m times m×n

Look at what the backward pass actually did. It started with a row vector at the loss and multiplied it through Jacobians, never building the big matrices. This is called a vector-Jacobian product. Because a loss is one number while parameters number in the millions, going backward gets every gradient in a single sweep. Going forward would need one sweep per parameter. That is why deep learning frameworks use reverse mode, with a cost of roughly a small multiple of the forward pass.

• Reverse mode: one pass, all gradients
• Vector-Jacobian products
• Cost near a few forward passes

Time to build the machine. We create a class called Value that wraps one number. It stores the data, a gradient that starts at zero, the parent nodes that produced it, and a small function that knows how to push gradients to those parents. Every arithmetic operation will create a new Value and remember how it came to be, quietly building the computational graph as you write ordinary math.

• data, grad
• parents, backward function

For addition, the new value out has parents self and other. Its backward function adds out dot grad to both parents' gradients, because the local derivative is one. For multiplication, self gets other dot data times out dot grad, and other gets self dot data times out dot grad. Notice the plus equals everywhere. That is fan-out accumulation, and using a plain equals sign is the number one bug in homemade autograd.

• Add: pass gradient to both
• Multiply: swap the other value
• Always use +=

A power operation with a fixed exponent n has local derivative n times x to the n minus one, multiplied by the upstream gradient. We do not need a new class for subtraction. We define it as self plus negative one times other, reusing add and multiply. Building big operations from tiny ones is exactly how real frameworks keep their set of primitives small.

• Power: n · x^(n−1)
• Subtraction: add the negative

The backward method does three things. First, it sorts all nodes in topological order, meaning every node appears after the nodes it depends on, using a depth-first search. Second, it sets the output gradient to one. Third, it walks through the sorted list in reverse and calls each node's backward function. Reverse order guarantees that a node's full gradient is complete before it passes anything further back.

• Topological order
• Seed grad = 1
• Call _backward in reverse

Now run our loss. Make Values for w equals two, x equals three, b equals one, y equals ten. Compute f as w times x plus b minus y, all squared. Printing f dot data gives nine. Call f dot backward. Out come w dot grad of negative eighteen, x dot grad of negative twelve, b dot grad of negative six, and y dot grad of six. These are exactly our hand results, produced by about forty lines of code.

• w=2, x=3, b=1, y=10
• f.data = 9
• w.grad = −18, b.grad = −6

Always test your autograd with the finite-difference check we used earlier. The two numbers should agree to many decimal places. Two pitfalls to remember. Because gradients accumulate, you must reset them to zero before each new backward pass in a training loop, or old gradients pile up. And fan-out works automatically: compute x times x plus three times x for x equal to two, and x dot grad comes out as seven.

• Compare with finite differences
• Reset grads between steps
• Reuse of nodes accumulates

Notice what we did not need. A neural network is just a larger computational graph, with matrix multiplications and nonlinearities as nodes. Each node still stores its inputs, knows its local derivative, and multiplies by the upstream gradient. The backpropagation in lesson 26 is this algorithm applied to layers, and PyTorch autograd in lesson 49 is a fast, tensor-valued version of the class we just wrote.

• Lesson 26: backprop in networks
• Lesson 49: PyTorch autograd

Let us recap. The chain rule says rates multiply along a chain: dy by dx equals dy by du times du by dx. With many variables, multiply along each path and add across paths. A computational graph stores forward values, then the backward pass multiplies each upstream gradient by a local derivative, accumulating wherever values fan out. In vector form, Jacobians multiply, and reverse mode computes cheap vector-Jacobian products.

• Rates multiply along a chain
• Sum over paths
• Upstream × local, accumulate with +=

In the next lesson we ask when gradient descent is guaranteed to find the best answer, through convexity, and how to optimize under constraints with Lagrange multipliers. Try the autograd class yourself, break it on purpose, and check it with finite differences. See you in lesson 37.

• Lesson 37: convexity
• Lagrange multipliers
• Subscribe for the full course