A3 · Advanced Mathematics for AI

The Chain Rule and Automatic Differentiation

The chain rule gives the derivative of a function inside another function: multiply the derivative of the outer function by the derivative of the inner one. If y depends on u and u depends on x, then dy/dx = dy/du · du/dx.

20 min Advanced

What you need: The earlier topics of this module, and its prerequisites.

Definition

The chain rule gives the derivative of a function inside another function: multiply the derivative of the outer function by the derivative of the inner one. If y depends on u and u depends on x, then dy/dx = dy/du · du/dx.

In simple words

Three gears are connected. Gear A turns gear B twice as fast, and gear B turns gear C three times as fast. So C turns 2 × 3 = 6 times as fast as A. Rates along a chain multiply. That's the whole rule.

Why it matters

A neural network is functions inside functions, layer upon layer. Backpropagation, the algorithm that trains every network, is the chain rule applied from the last layer back to the first. Understand this one rule and backpropagation stops being mysterious.

How it works

  1. Spot the outer and inner functions. In (3x + 1)², the outer is "square it" and the inner is 3x + 1.
  2. Differentiate the outer, leaving the inner alone: 2(3x + 1).
  3. Multiply by the derivative of the inner: × 3.

Why it matters for a network

A tiny network computes hidden = w₁·x, then output = w₂·hidden, then loss = (output − y)². To find how the first weight affects the loss, multiply the rates along the chain from the loss back to that weight. Backpropagation does exactly this, and it reuses the shared pieces so that even billions of weights can be handled.

The maths

Chain rule
The inner function, u = g(x)
The outer function, y = f(u)
A longer chain

Numbers for the tiny network: x = 2, w₁ = 3, w₂ = 0.5, y = 1.

Forward: h = 3 × 2 = 6     ŷ = 0.5 × 6 = 3     L = (3 − 1)² = 4

∂L/∂ŷ = 2(ŷ − y) = 2 × 2 = 4
∂ŷ/∂h = w₂ = 0.5
∂h/∂w₁ = x = 2
∂L/∂w₁ = 4 × 0.5 × 2 = 4

The result is 4, a positive slope, so reducing w₁ reduces the loss.

In code

chain_rule.py
x, y = 2, 1
w1, w2 = 3.0, 0.5

def loss(w1):
    return (w2 * (w1 * x) - y) ** 2

# Chain rule by hand
h = w1 * x
y_hat = w2 * h
by_chain_rule = 2 * (y_hat - y) * w2 * x

# Check it numerically
step = 1e-6
numeric = (loss(w1 + step) - loss(w1)) / step

print("Loss:", loss(w1))
print("Chain rule:", by_chain_rule)
print("Numeric check:", round(numeric, 4))

Output

Loss: 4.0
Chain rule: 4.0
Numeric check: 4.0

Common mistakes

  • Forgetting to multiply by the inner derivative. The derivative of (3x+1)² is 6(3x+1), not 2(3x+1).
  • Stopping halfway along a long chain. Every link in the chain must be multiplied.

Check yourself

Question 1

1. What is the derivative of (2x + 5)³?

Question 2

2. Which algorithm is the chain rule applied across a network?

Summary

  • For nested functions, derivatives multiply.
  • Outer's derivative × inner's derivative.
  • Backpropagation is this rule, run from the loss back to every weight.

Finished reading?

Sources: Deisenroth, Faisal & Ong, Mathematics for Machine Learning