🔗 Chain rule for composed functions

By Better Wise

Why compositions need their own rule

You already have the two ingredients: the gradient vector of a scalar function, and the Jacobian of a vector-valued function. Almost nothing you differentiate in practice arrives expanded, because most functions are compositions: a linear map feeding an elementwise activation, a log wrapping an exponential, a sum of squares wrapping an affine map. The chain rule is the one rule that handles all of them without expanding anything.

After this lesson you can:

  • name the intermediate variables of a composition and write the forward pass;
  • assemble the Jacobian product in the correct order;
  • collapse that product to a gradient when the outer map is scalar;
  • confirm every hand-derived gradient with a central finite difference in Python.

One-line recall from the previous two lessons: the gradient of $f:\mathbb{R}^n \to \mathbb{R}$ is the column vector of partials; the Jacobian of $h:\mathbb{R}^n \to \mathbb{R}^m$ is the $m \times n$ matrix whose row $i$ is $\nabla h_i$.

The chain rule in Jacobian form

Write the composition as a chain of named maps, $f = g \circ h$, so $x \mapsto h(x) = y \mapsto g(y) = f(x)$. With $h:\mathbb{R}^n \to \mathbb{R}^m$ and $g:\mathbb{R}^m \to \mathbb{R}^k$, the derivative of the composition is the product of derivatives evaluated at matching points:

$$ J_f(x) = J_g(h(x)), J_h(x). $$

Shapes confirm the order: $(k \times m)(m \times n) = k \times n$. The inner Jacobian is evaluated at $x$; the outer Jacobian is evaluated at _____. Reversing the order gives a shape mismatch, which is why the ordering mistake is usually caught by arithmetic rather than by luck.

flowchart TD
    X["🧮<br/>x"] --> H["🔁<br/>h(x) = y"]
    H --> G["🎯<br/>g(y) = f(x)"]
    JX["J_h(x)"] --> P["J_g(h(x)) J_h(x)"]
    JY["J_g(h(x))"] --> P

When $k = 1$, $J_g$ is a single row, the row vector $\nabla g(y)^\top$, and the product becomes a row vector. Transposing gives the gradient of the composition, which is worth memorizing in its own right:

$$ \nabla f(x) = J_h(x)^\top \nabla g(h(x)). $$

In this form the cloze reads: the gradient of a scalar-valued composition is _____, or the inner Jacobian transposed times the outer gradient.

Worked example: a nested scalar composition

Let $f(x) = e^{x_1^2 + 3x_2}$, so $h(x) = x_1^2 + 3x_2$ and $g(u) = e^u$. Here $J_h(x) = [2x_1,; 3]$ and $g'(u) = e^u$, so

$$ \nabla f(x) = e^{x_1^2+3x_2}\begin{bmatrix} 2x_1 \ 3 \end{bmatrix}. $$

Now wrap a log around it: $L = \log f$. Then $L = x_1^2 + 3x_2$ and $\nabla L = [2x_1, 3]^\top$. The scalar factors $f$ and $1/f$ cancel, which is the chain rule telling you composition can simplify, not complicate.

Worked example: composing through a matrix-shaped Jacobian

Take $h:\mathbb{R}^2 \to \mathbb{R}^2$ with $h(x) = (x_1 x_2,; x_1 + x_2^2)$ and $g(y) = y_1 y_2$.

Hand derivation:

$$ J_h(x) = \begin{bmatrix} x_2 & x_1 \ 1 & 2x_2 \end{bmatrix}, \qquad \nabla g(y) = \begin{bmatrix} y_2 \ y_1 \end{bmatrix}. $$

Multiplying $J_h(x)^\top \nabla g(y)$ and substituting $y = h(x)$:

$$ \nabla f(x) = \begin{bmatrix} x_2 y_2 + y_1 \ x_1 y_2 + 2x_2 y_1 \end{bmatrix} = \begin{bmatrix} 2x_1x_2 + x_2^3 \ x_1^2 + 3x_1x_2^2 \end{bmatrix}. $$

Direct expansion gives $f = x_1^2 x_2 + x_1 x_2^3$, and differentiating that term by term lands on exactly the same vector. Two routes, one answer.

Worked example: elementwise layers produce diagonal Jacobians

Compositions that apply a scalar function to each coordinate are everywhere. For $a = Wx + b$ with $W \in \mathbb{R}^{2 \times 2}$, $t = \tanh(a)$ (elementwise), and readout $f = r^\top t$:

$$ J_a(x) = W, \qquad J_t(a) = \operatorname{diag}(1 - t_1^2,, 1 - t_2^2), \qquad \nabla_t f = r. $$

Chaining: $\nabla_x f = W^\top \operatorname{diag}(1-t^2), r = W^\top\big((1 - t^2) \odot r\big)$. An elementwise map always contributes a _____ Jacobian, so the matrix product degrades into an elementwise scaling. That is the whole reason backpropagation is cheap.

Confirming the chain with finite differences

Never ship a hand-derived gradient without a numeric check. The central difference for coordinate $i$ is _____, which is second-order accurate: its truncation error is $O(h^2)$, while a one-sided difference is only $O(h)$.

import numpy as np

def f(x):                                       # g o h, R^2 -> R
    y1, y2 = x[0] * x[1], x[0] + x[1]**2
    return y1 * y2

def grad_analytic(x):
    x1, x2 = x
    return np.array([2*x1*x2 + x2**3, x1**2 + 3*x1*x2**2])

def grad_fd(fun, x, h=1e-6):
    g = np.zeros_like(x)
    for i in range(x.size):
        e = np.zeros_like(x); e[i] = h
        g[i] = (fun(x + e) - fun(x - e)) / (2 * h)
    return g

x = np.array([1.3, -0.7])
ga, gn = grad_analytic(x), grad_fd(f, x)
print(ga, gn, np.linalg.norm(ga - gn) / np.linalg.norm(gn))

The relative error comes out around $10^{-10}$. Anything below $10^{-6}$ is a pass; a relative error near $10^{-2}$ means a missing factor, a wrong transpose, or the outer Jacobian evaluated at $x$ instead of at $h(x)$.

Practical rules for chaining derivatives

  • Name your intermediates. Write $a = Wx + b$, $t = \tanh(a)$, $f = r^\top t$ before touching a matrix. Most chain-rule errors are bookkeeping errors.
  • Check shapes first. If a product does not have the shape of the answer, the order is wrong. Do this before checking values.
  • Mind the evaluation points. Inner Jacobian at the input, outer Jacobian at the intermediate value. This is the single most common bug in hand derivations.
  • Elementwise means diagonal. Never build the diagonal explicitly in code; write it as an elementwise product.
  • Use $h \approx 10^{-6}$ for central differences in float64. Smaller $h$ lets roundoff dominate; larger $h$ lets truncation dominate.
  • Test at a random point, not a tidy one. Zeros and ones hide sign and factor errors.

What to carry forward

The chain rule says the derivative of a composition is the ordered product of the derivatives of its parts, with the gradient case being $J_h(x)^\top \nabla g(h(x))$; everything else in this lesson is bookkeeping around that one line, checked numerically.

Next comes the checkpoint on first-order derivatives, where you will be asked to move between gradients, Jacobians, and chained compositions on demand.

Front: Jacobian of a composition f = g ∘ h Back: J_f(x) = J_g(h(x)) J_h(x), outer times inner.

Front: Gradient of a scalar-valued composition f = g ∘ h Back: ∇f(x) = J_h(x)^T ∇g(h(x)).

Front: At which point is the outer Jacobian J_g evaluated? Back: At the intermediate value h(x), not at x.

Front: At which point is the inner Jacobian J_h evaluated? Back: At the input x.

Front: Jacobian of an elementwise function t = φ(a) Back: A diagonal matrix diag(φ'(a_i)).

Front: Jacobian of the affine map a = Wx + b Back: Just W, independent of x.

Front: Gradient of f(x) = e^(x1² + 3x2) Back: e^(x1² + 3x2) times [2x1, 3].

Front: Central difference formula for coordinate i Back: (f(x + h e_i) - f(x - h e_i)) / (2h).

Front: Truncation error order of a central difference Back: O(h²), versus O(h) for a one-sided difference.

Front: Why not use h = 1e-14 in a finite-difference check? Back: Roundoff in the subtraction swamps the difference; 1e-6 is a good float64 default.

Front: Gradient of f = r^T tanh(Wx + b) Back: W^T ((1 - tanh(Wx+b)²) ⊙ r).

Front: Shape test that catches reversed chain-rule products Back: (k×m)(m×n) must equal the k×n Jacobian of the composition.

Front: What does a relative error near 1e-2 in a finite-difference check usually mean? Back: A missing factor, a wrong transpose, or the outer Jacobian evaluated at the wrong point. A last recall sweep before the practical rules.

The multivariate chain rule for Jacobians reads J_f(x) = _____, outer times inner. For a scalar-valued composition f = g ∘ h the gradient is ∇f(x) = _____. An elementwise nonlinearity t = φ(a) has a _____ Jacobian, diag(φ'(a_i)). A central difference has truncation error _____, one order better than a one-sided difference. A reliable default step for a float64 finite-difference check is _____.

Rules of thumb that save you from silent gradient bugs

  • Write the shapes before the algebra. Label every intermediate with its dimension, for example x is n, a is m, f is a scalar. If the product you wrote cannot be formed, the chain rule is written in the wrong order, not the wrong function.
  • The outer Jacobian never sees x. A very common slip is writing J_g(x) J_h(x) for f = g ∘ h. The outer factor lives at the intermediate value h(x). Numerically this produces a plausible-looking but wrong gradient, which is why the finite-difference check matters.
  • Transpose only when you leave Jacobian land. Keep everything as Jacobians while chaining, then apply the transpose once, at the end, if you want a gradient rather than a row vector. Transposing mid-chain is the second most common error.
  • Do not shrink h out of enthusiasm. h = 1e-12 in float64 gives garbage because the two function values agree to twelve digits and the subtraction loses them. Start at 1e-6, and if the error looks large, try 1e-5 and 1e-7 to see whether roundoff or truncation dominates.
  • Judge with relative error, not absolute error. Compare with np.linalg.norm(num - ana) / (np.linalg.norm(num) + np.linalg.norm(ana)). Values near 1e-7 mean agreement; values near 1e-2 mean a missing factor, a wrong transpose, or the outer Jacobian at the wrong point.
  • Reuse one checker. A single function that takes a callable and its analytic gradient and loops over coordinates will catch more bugs over the course than any amount of staring at algebra.
import numpy as np

def check_grad(f, grad_f, x, h=1e-6):
    x = np.asarray(x, dtype=float)
    num = np.zeros_like(x)
    for i in range(x.size):
        e = np.zeros_like(x); e[i] = h
        num[i] = (f(x + e) - f(x - e)) / (2 * h)
    ana = np.asarray(grad_f(x), dtype=float)
    rel = np.linalg.norm(num - ana) / (np.linalg.norm(num) + np.linalg.norm(ana) + 1e-12)
    return rel, num, ana

What to carry forward

You now have the full first-order toolkit: partials assemble into a gradient, vector-valued maps get a Jacobian, and compositions multiply those Jacobians from the outside in, with the outer factor evaluated at the intermediate value. The reusable habit is the pairing: derive by hand, then confirm with a central difference at 1e-6 and a relative error near 1e-7.

The next lesson is the module checkpoint, where you will derive a few gradients and Jacobians from scratch, chain them, and defend each answer with a numeric check, with no new theory introduced.

Front: Chain rule for f = g ∘ h, in Jacobian form Back: J_f(x) = J_g(h(x)) J_h(x), outer times inner.

Front: Chain rule for a scalar-valued composition, in gradient form Back: ∇f(x) = J_h(x)^T ∇g(h(x)).

Front: Where is the outer Jacobian evaluated? Back: At h(x), the intermediate value, never at x.

Front: Jacobian of an elementwise map t = φ(a) Back: The diagonal matrix diag(φ'(a_i)).

Front: Jacobian of the affine map a = Wx + b Back: W, and it does not depend on x.

Front: Gradient of f(x) = exp(x1² + 3x2) Back: exp(x1² + 3x2) · [2x1, 3].

Front: Gradient of f = rᵀ tanh(Wx + b) Back: Wᵀ ((1 − tanh(Wx+b)²) ⊙ r).

Front: Central difference for coordinate i Back: (f(x + h e_i) − f(x − h e_i)) / (2h).

Front: Truncation order of central vs one-sided differences Back: O(h²) for central, O(h) for one-sided.

Front: Why avoid h = 1e-14 in float64? Back: Roundoff in the subtraction swamps the difference; 1e-6 is a sound default.

Front: Shape test for a chained product Back: (k×m)(m×n) must give the k×n Jacobian of the composition.

Front: Relative error near 1e-2 in a gradient check usually means Back: A missing factor, a wrong transpose, or the outer Jacobian evaluated at the wrong point.