Chain Rule

How This Lesson Fits the Module

Derivatives taught how to measure instantaneous change for a single variable. Partial Derivatives extended that idea to functions of many variables—holding all but one input fixed while measuring sensitivity.

Real ML models are almost never “one function of one variable.” They are compositions: layers stacked on layers, activations wrapped around linear transforms, loss built from predictions and labels. The chain rule is the calculus tool that tells us how change propagates through those compositions.

Every automatic differentiation engine—PyTorch’s autograd, TensorFlow’s gradient tape, JAX’s grad—implements the chain rule at scale. Understanding it turns backpropagation from a black-box recipe into a principled algorithm you can debug, extend, and reason about.

Learning Objectives

By the end of this lesson, students should be able to:

State and apply the single-variable chain rule for compositions f(g(x)).
Extend the chain rule to multivariable settings: scalar outputs, vector intermediates, and multiple paths.
Write the chain rule in summation and Jacobian form and recognize when each view is clearer.
Draw and traverse a computational graph to propagate local derivatives backward.
Explain backpropagation as repeated application of the multivariable chain rule.
Compute gradients for small neural-network-style expressions by hand.
Connect manual chain-rule reasoning to framework behavior (.backward(), grad).
Identify common chain-rule mistakes: forgotten inner derivatives, path confusion, and shape mismatches.

Introduction: Why Composition Demands a New Rule

Suppose a model computes a prediction in two stages. First an intermediate value u = g(x), then a loss contribution L = f(u). You know ∂f/∂u (how L responds to u) and dg/dx (how u responds to x). What is dL/dx?

You cannot simply multiply the functions f and g and differentiate—that would ignore the nested structure. Nor can you add their derivatives. The correct answer multiplies their rates of change: a small change in x perturbs u, which perturbs L. The chain rule quantifies that two-step sensitivity as a product.

In deep learning, “two steps” becomes millions. The same principle scales: local derivatives at each node, multiplied along paths, summed when paths merge.

The Single-Variable Chain Rule

Theorem — Chain Rule (One Variable)

If g is differentiable at x and f is differentiable at g(x), then the composition h(x) = f(g(x)) is differentiable at x and

h′(x) = f′(g(x)) · g′(x)

Equivalent Leibniz notation: if y = f(u) and u = g(x), then

dy/dx = (dy/du) · (du/dx)

Intuition. A small change Δx produces a change in u of roughly (du/dx)Δx. That change in u produces a change in y of roughly (dy/du) times the u-perturbation. Multiply the two rates to get the overall rate dy/dx.

Worked Example 1 — Polynomial Composition

Problem. Let h(x) = (x² + 1)³. Find h′(x).

Setup. Define u = g(x) = x² + 1 and y = f(u) = u³.

Local derivatives. du/dx = 2x, and dy/du = 3u² = 3(x² + 1)².

Chain rule. h′(x) = 3(x² + 1)² · 2x = 6x(x² + 1)².

Check. Expanding with the power rule directly gives the same result—the chain rule is a structured shortcut, not a different answer.

Worked Example 2 — Exponential of a Linear Function

Problem. Differentiate h(x) = e^2x+1.

Let u = 2x + 1. Then du/dx = 2 and d/du[e^u] = e^u.

Therefore h′(x) = e^2x+1 · 2 = 2e^2x+1.

ML connection. Sigmoid and softmax involve exponentials of affine expressions. The inner derivative (here, 2) is the weight vector contribution in a logit layer.

Longer Chains

For y = f(g(h(x))), apply the rule repeatedly:

dy/dx = (dy/du) · (du/dv) · (dv/dx)

Each link multiplies. A 50-layer network is a 50-factor product of local Jacobians (matrices in the multivariable case). Vanishing or exploding gradients arise when many of those factors are consistently smaller or larger than 1.

From One Variable to Many: The Multivariable Chain Rule

ML functions rarely map ℝ → ℝ. A layer maps a weight vector to activations; a loss maps predictions and labels to a scalar. We need the chain rule when:

The output is a scalar (loss) but inputs are vectors (weights).
Intermediate variables are vectors (hidden activations).
One variable influences the output along multiple paths (skip connections, shared weights).

Theorem — Chain Rule (Scalar Output)

Let z = f(y₁, …, y_m) be a scalar function of intermediate variables, and let each y_j depend on inputs x₁, …, x_n. Then for each input x_i:

∂z/∂x_i = ∑_j=1^m (∂z/∂y_j) (∂y_j/∂x_i)

Read this as: sum over paths through each intermediate y_j the product of (sensitivity of output to intermediate) × (sensitivity of intermediate to input).

PrerequisitePartial derivatives ∂z/∂y_j and ∂y_j/∂x_i are defined in Partial Derivatives. The chain rule tells us how to combine them.

Vector Form and the Jacobian

When intermediates are vectors, compact notation uses the Jacobian matrix. If y = g(x) with x ∈ ℝⁿ, y ∈ ℝ^m, and scalar L = f(y), then:

∇_xL = J_g(x)^T ∇_yL

Here J_g is the m × n matrix of partial derivatives (∂y_i/∂x_j), and the gradient ∇_yL is a column vector. The transpose appears because we accumulate contributions into each component of x. This single line is the engine of reverse-mode automatic differentiation.

Summation Form

Best for hand calculations and small graphs. Explicitly lists each path through intermediate nodes.

∂L/∂x = ∑_j (∂L/∂y_j)(∂y_j/∂x)

Jacobian–Vector Product

Best for implementation. Frameworks never materialize full Jacobians for huge layers; they multiply by gradient vectors efficiently (vector–Jacobian products).

∇_xL = J^T ∇_yL

Worked Example 3 — Two Paths to the Same Input

Problem. Let x, y ∈ ℝ and define u = xy, v = x + y, L = u² + v². Find ∂L/∂x.

Local partials. ∂L/∂u = 2u = 2xy, ∂L/∂v = 2v = 2(x + y). Also ∂u/∂x = y, ∂v/∂x = 1.

Chain rule (sum over paths).

∂L/∂x = (∂L/∂u)(∂u/∂x) + (∂L/∂v)(∂v/∂x)

= (2xy)(y) + (2(x + y))(1) = 2xy² + 2x + 2y

Lesson. When x affects L through more than one intermediate, add the contributions. This is the multivariable generalization of “don’t forget a path.”

Computational Graphs: Visualizing Composition

A computational graph is a directed acyclic graph (DAG) whose nodes are variables or operations and whose edges show data flow. Each edge carries a forward value; each node stores a rule for local differentiation.

Forward pass — Evaluate nodes in topological order: inputs → intermediates → loss ↓ Store activations — Save values needed for backward (memory–compute trade-off) ↓ Backward pass — Seed ∂L/∂L = 1, propagate gradients upstream via chain rule ↓ Accumulate at forks — Sum gradients when one variable feeds multiple downstream nodes

Consider a tiny computation: x → square → u = x², and u → sin → y = sin(u). The graph is a chain of two nodes after the input.

Node	Forward value (example x = 2)	Local derivative	Backward message
u = x²	u = 4	du/dx = 2x	(∂L/∂u) · 2x
y = sin(u)	y = sin(4)	dy/du = cos(u)	∂L/∂u = (∂L/∂y) · cos(u)

For L = y, starting from ∂L/∂y = 1: ∂L/∂u = cos(4), then ∂L/∂x = cos(4) · 4. This matches the single-variable chain rule d/dx[sin(x²)] = cos(x²) · 2x.

Engineering Principle

Frameworks build this graph implicitly from tensor operations. When you call loss.backward(), the runtime walks the graph in reverse, applying the same local-derivative multiplication you would do by hand—at billion-node scale.

Backpropagation Is the Chain Rule

Backpropagation is not a separate invention from calculus. It is reverse-mode automatic differentiation: apply the multivariable chain rule once per edge, from loss back to parameters, reusing intermediate results.

Definition — Backpropagation (Informal)

Given a scalar loss L computed by a differentiable program (the network), backpropagation computes ∂L/∂θ for every parameter θ by:

Running the forward pass to evaluate L and store activations.
Initializing the output adjoint ¯L = ∂L/∂L = 1.
Visiting nodes in reverse topological order, multiplying each node’s local Jacobian (or scalar derivative) by the incoming adjoint and passing results to parent nodes.
Summing adjoints at nodes with multiple children (forks in the graph).

Worked Example 4 — One Neuron with Sigmoid

Setup. A single neuron with weight w, bias b, input x:

z = wx + b, a = σ(z) = 1/(1 + e^−z), L = ½(a − y)²

where y is a target label. Compute ∂L/∂w.

Forward (concrete numbers). Take x = 2, w = 0.5, b = −1, y = 1. Then z = 0, a = 0.5, L = 0.125.

Backward (chain rule step by step).

∂L/∂a = a − y = 0.5 − 1 = −0.5
∂a/∂z = σ(z)(1 − σ(z)) = 0.25 (sigmoid derivative at z = 0)
∂L/∂z = (∂L/∂a)(∂a/∂z) = (−0.5)(0.25) = −0.125
∂z/∂w = x = 2
∂L/∂w = (∂L/∂z)(∂z/∂w) = (−0.125)(2) = −0.25

Similarly ∂L/∂b = ∂L/∂z = −0.125. Gradient descent would update w ← w − η(∂L/∂w) to reduce loss.

Worked Example 5 — Affine Layer + ReLU

Setup. z = Wx + b, a = ReLU(z) (element-wise), scalar L = ∑_i a_i (simplified loss for illustration).

Backward. ∂L/∂a_i = 1 for all i. ReLU gate: ∂a_i/∂z_i = 1 if z_i > 0, else 0. So ∂L/∂z_i = 1 on active units, 0 on inactive.

For weights: z_i = ∑_j W_ijx_j + b_i, hence ∂z_i/∂W_ij = x_j and

∂L/∂W_ij = (∂L/∂z_i) · x_j

This outer-product structure (“adjoint times forward input”) appears in every linear layer and is why weight gradients are implemented as matrix multiplications, not naive loops.

Forward Mode vs Reverse Mode

The chain rule can be traversed in two directions. ML training almost always uses reverse mode because we have one scalar loss and millions of parameters.

Mode	Direction	Cost (rough)	Typical Use
Forward mode	Propagates ∂(·)/∂x alongside forward values for one input direction	One forward pass per input dimension	Jacobian–vector products when n ≪ m
Reverse mode	Propagates ∂L/∂(·) backward from the loss	One backward pass for all parameters	Training neural networks (n huge, one scalar loss)

Coming UpThe gradient vector ∇f packages all first partial derivatives into one object. Gradient formalizes that packaging and its geometric meaning as the direction of steepest ascent.

Chain Rule in Code

Manual backprop is error-prone; frameworks automate the graph traversal. Still, reading autograd output should match hand-derived gradients.

import torch

x = torch.tensor(2.0, requires_grad=True)
u = x ** 2
y = torch.sin(u)
L = y

L.backward()          # reverse-mode chain rule
print(x.grad)         # cos(4) * 4  ≈  -2.614

# Tiny neuron: L = 0.5 * (sigmoid(w*x + b) - y)**2
w = torch.tensor(0.5, requires_grad=True)
b = torch.tensor(-1.0, requires_grad=True)
x = torch.tensor(2.0)
y_true = torch.tensor(1.0)

z = w * x + b
a = torch.sigmoid(z)
loss = 0.5 * (a - y_true) ** 2
loss.backward()

print(w.grad)         # ≈ -0.25  (matches Worked Example 4)
print(b.grad)         # ≈ -0.125

PyTorch records operations on tensors with requires_grad=True, building a dynamic computational graph. backward() applies the multivariable chain rule exactly as derived above. TensorFlow and JAX follow the same mathematical pattern with different APIs.

Common Misconceptions

Misconception 1: “Backpropagation is a heuristic separate from calculus.”

Why people believe it: Historical naming and engineering-focused tutorials present update rules without derivatives.

Reality: Backprop is reverse-mode AD, which is the multivariable chain rule organized for efficient computation. Every gradient it produces is a partial derivative.

Misconception 2: “You can forget the inner derivative if the outer function is simple.”

Why people believe it: Memorized derivative formulas for outer functions (sin, exp, ReLU) without tracking the argument.

Reality: d/dx sin(x²) is cos(x²) · 2x. Omitting the inner factor is the most common hand-calculation bug and the analog of broken gradients in custom layers.

Misconception 3: “At a fork in the graph, multiply the branch gradients.”

Why people believe it: Confusion between chain (multiply along a path) and fork (split) semantics.

Reality: Multiply along each path; add when several paths share the same input. Worked Example 3 demonstrates this explicitly.

Misconception 4: “Frameworks compute the full Jacobian for every layer.”

Why people believe it: Jacobian notation in textbooks suggests huge matrices are formed.

Reality: Reverse mode computes vector–Jacobian products ∇_yL ↦ J^T∇_yL without storing J explicitly. That efficiency is what makes training large models feasible.

Quick Knowledge Check

Computation: Differentiate (3x − 1)⁴. Answer: 4(3x − 1)³ · 3 = 12(3x − 1)³.
True/False: For y = f(g(x)), the chain rule gives y′ = f′(x) · g′(x). Answer: False — evaluate f′ at g(x), not at x.
Short Answer: What does backpropagation compute? Answer: Partial derivatives of a scalar loss with respect to parameters (gradients), via reverse-mode chain rule.
Computation: If L = u², u = 2x, find dL/dx at x = 3. Answer: (2u)(2) = 4u = 24.
Multiple Choice: At a graph fork where one variable feeds two nodes, gradients: (a) multiply, (b) add, (c) average, (d) take the max. Answer: (b) add.
Short Answer: Why is reverse mode preferred when L is scalar and parameters are millions of dimensions? Answer: One backward pass yields all ∂L/∂θ_i; forward mode would need one pass per parameter.
Computation: For σ(z) = 1/(1+e^−z), what is σ′(z) at z = 0? Answer: 0.25.
True/False: The computational graph for training must be a DAG (no cycles). Answer: True for standard feedforward backprop; RNNs unroll in time to form a DAG per step.
Short Answer: In ∂L/∂W_ij = (∂L/∂z_i)x_j, what role does x_j play? Answer: Forward input to that connection—the inner partial ∂z_i/∂W_ij.
Multiple Choice: d/dx e^x² equals: (a) e^x², (b) 2xe^x², (c) e^2x, (d) 2e^x². Answer: (b).

Key Takeaways

The chain rule multiplies local rates of change along composed functions: (df/dg) · (dg/dx).
With multiple variables, sum over paths: ∂z/∂x_i = ∑_j (∂z/∂y_j)(∂y_j/∂x_i).
Computational graphs make composition explicit; backprop walks them in reverse.
Backpropagation is reverse-mode automatic differentiation—not a separate trick from calculus.
Neural network layers are chain-rule instances: activation derivatives times upstream adjoints, weight gradients from adjoint × forward inputs.
Frameworks implement vector–Jacobian products efficiently; they do not materialize full Jacobians.
Mastering the chain rule bridges Partial Derivatives to the Gradient and training algorithms that follow.