The Chain Rule

A neural network is a composition of functions, so its derivative is a product of local derivatives. The chain rule is that product. It is the single fact that makes backpropagation possible, and it is worth seeing both the scalar form and the sum-over-paths form.

The idea

For $y = f(u)$ with $u = g(x)$:

The intuition: a change in $x$ changes $u$ at rate $du/dx$, and that change in $u$ changes $y$ at rate $dy/du$. Multiply the rates.

When one variable feeds several paths, add the contributions. If $z$ depends on intermediate variables $u_1, \dots, u_k$, each of which depends on $x$:

This total-derivative form is what a computational graph encodes. Reverse-mode autodiff walks the graph from the output backward, and at each node multiplies the incoming gradient by that node's local derivative and adds it to the parent's accumulator. One backward pass yields the gradient with respect to every parameter, at roughly the cost of one forward pass.

Worked example

Differentiate $y = (3x + 1)^2$. Set $u = 3x + 1$. Then

At $x = 2$ that is $42$. Expanding first confirms it: $y = 9x^2 + 6x + 1$, and $y' = 18x + 6 = 42$.

The sum form shows up with a shared node. If $z = u^2 + 3u$ and $u = 2x$, then $dz/du = 2u + 3 = 4x + 3$ and $du/dx = 2$, giving $dz/dx = 8x + 6$. The two terms of $dz/du$ are two paths through the graph, and both get multiplied by the same $du/dx$.

In code

import torch

x = torch.tensor(2.0, requires_grad=True)
u = 3 * x + 1
y = u**2
y.backward()
print(u.item())    # 7.0
print(x.grad)      # tensor(42.)

Autograd stored $u$ during the forward pass so the backward pass could reuse it. That storage is why training costs memory proportional to depth: the graph holds the intermediates.

Check yourself

1. Differentiate $e^{x^2}$ and $\sin(2x)$ using the chain rule.
2. Why does backward mode need the forward intermediates, and what does that imply for memory?
3. Use the chain rule on $z = uv$ with $u = u(x)$ and $v = v(x)$ to recover the product rule.

Key takeaways

• The chain rule multiplies local derivatives along a path and sums across paths.
• Reverse mode computes all parameter gradients in one backward pass.
• Backprop is the chain rule on a computational graph, not a special algorithm.