Multivariable Chain Rule

Why you need a new-looking chain rule

In 1D, composition looks like $h(x) = f(g(x))$ and the chain rule says $h'(x) = f'(g(x))\, g'(x)$.

In multiple dimensions, you still compose functions, but now each function can take multiple inputs and produce multiple outputs. The “derivative” is no longer a single number; it’s a linear map that best approximates the function near a point. Linear maps are represented by matrices, so “multiply the derivatives” becomes “multiply the Jacobian matrices.”

This is the conceptual core:

• •Derivative = best linear approximation near a point.
• •Jacobian matrix = the matrix of partial derivatives that represents that linear approximation.
• •Composition of linear maps corresponds to matrix multiplication.

The composition structure (track shapes)

The cleanest multivariable chain rule is written with a clear inner/outer structure:

• •Inner map: $g: \mathbb{R}^n \to \mathbb{R}^m$
• •Outer map: $f: \mathbb{R}^m \to \mathbb{R}^p$
• •Composition: $h = f \circ g: \mathbb{R}^n \to \mathbb{R}^p$

Let:

• •input vector x ∈ ℝⁿ
• •intermediate vector u = g(x) ∈ ℝᵐ
• •output vector y = f(u) ∈ ℝᵖ

A quick “shape table” you should get used to:

Notice the only multiplication that makes sense dimensionally is:

That is already a big part of the multivariable chain rule: the shapes force the correct order.

The chain rule (matrix form)

The multivariable chain rule states:

Read it as: “first apply the derivative of $g$ at x, then apply the derivative of $f$ at the resulting point $g(\mathbf{x})$.”

What it means geometrically (tiny perturbations)

If you perturb the input by a small vector $\Delta \mathbf{x}$, then:

1) Push forward through g:

2) Push forward through f:

Combine them:

So the Jacobian of the composition is the matrix that takes $\Delta \mathbf{x}$ directly to $\Delta \mathbf{y}$. That matrix is the product $Df\,Dg$.

A note on notation: Df vs ∇f

You said you know gradients already. The key is to keep these distinct:

• •$Df$ is the Jacobian (a matrix, in general).
• •$\nabla f$ is the gradient (a vector) and is defined when $f$ is scalar-valued: $f: \mathbb{R}^n \to \mathbb{R}$.

When $p=1$, the Jacobian of $f$ is a 1×m row vector; the gradient is usually written as an m×1 column vector. They are transposes of each other (depending on convention):

This transpose issue matters a lot in backprop, so we’ll be explicit about it later.

Why this viewpoint

If you treat multivariable derivatives as “a bunch of partial derivatives,” you can still compute things, but it’s easy to lose track of structure.

If you treat the derivative as “the best linear map near a point,” everything becomes systematic:

• •You can push forward small changes.
• •You can compose derivatives by composing linear maps.
• •You can check correctness by verifying matrix shapes.

The linear approximation definition

Let $g: \mathbb{R}^n \to \mathbb{R}^m$. We say $g$ is differentiable at x if there exists a linear map $L: \mathbb{R}^n \to \mathbb{R}^m$ such that

with an error that becomes negligible compared to $\|\Delta \mathbf{x}\|$ as $\Delta \mathbf{x} \to 0$.

That linear map is the derivative $Dg(\mathbf{x})$.

Jacobian matrix: how the linear map is represented

Write $g(\mathbf{x}) = \begin{bmatrix} g_1(\mathbf{x}) \\ \vdots \\ g_m(\mathbf{x}) \end{bmatrix}$.

Then the Jacobian $Dg(\mathbf{x})$ is the m×n matrix:

This is exactly the matrix that maps a small input perturbation to an approximate output perturbation:

Interactive-canvas mental model (shape-tracking + arrows)

Imagine an “interactive canvas” with three boxes:

1) x-box: x ∈ ℝⁿ

2) u-box: u = g(x) ∈ ℝᵐ

3) y-box: y = f(u) ∈ ℝᵖ

Now add two kinds of arrows:

Arrow type A: pushforward of perturbations (forward mode)

You draw a little arrow $\Delta \mathbf{x}$ at the input.

• •It transforms to $\Delta \mathbf{u} = Dg\, \Delta \mathbf{x}$.
• •Then transforms to $\Delta \mathbf{y} = Df\, \Delta \mathbf{u}$.

So perturbations flow with the function direction.

Arrow type B: pullback of sensitivities / gradients (reverse mode)

If the final output is scalar (p=1), you draw a gradient arrow at the output: $\nabla_{\mathbf{y}} \ell$ (often just 1 if the scalar is the loss itself).

Then gradients flow backwards through transposes:

• •sensitivity at u: $\nabla_{\mathbf{u}} \ell = Df(\mathbf{u})^\top \nabla_{\mathbf{y}} \ell$
• •sensitivity at x: $\nabla_{\mathbf{x}} \ell = Dg(\mathbf{x})^\top \nabla_{\mathbf{u}} \ell$

This is the heart of backprop. In this lesson, we’re building the “transpose reflex”:

• •Perturbations push forward with $J$.
• •Gradients pull back with $J^\top$.

A small but crucial convention check

There are two common conventions:

1) Jacobian is m×n (outputs by inputs). Gradients are column vectors.

2) Jacobian is n×m (inputs by outputs). Gradients are row vectors.

We’ll use the most common ML convention:

• •Jacobian $Dg$ is m×n.
• •Gradient $\nabla f$ is n×1 for scalar f.

With that convention, the chain rule for scalar output becomes a clean transpose pullback (we’ll derive it soon).

Why matrix multiplication appears

The multivariable chain rule is not a new rule you memorize. It is a consequence of one fact:

If you approximate each function by a linear map near the relevant point, then approximating the composition means composing those linear maps.

And linear maps compose by matrix multiplication.

Derivation (showing the work)

Let $g: \mathbb{R}^n \to \mathbb{R}^m$ and $f: \mathbb{R}^m \to \mathbb{R}^p$.

Define:

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

Start with a small perturbation $\Delta \mathbf{x}$.

Step 1: Linearize g at x

Let $\Delta \mathbf{u} = Dg(\mathbf{x})\,\Delta\mathbf{x}$, so

Step 2: Linearize f at u

Substitute $\Delta \mathbf{u} = Dg(\mathbf{x})\,\Delta\mathbf{x}$:

But the left side is approximately

So we have the linear approximation:

By uniqueness of the best linear approximation, the derivative must be:

That’s the multivariable chain rule.

Element-wise chain rule (path-tracing through indices)

Sometimes you want the version that looks like “sum over paths.”

Let $h = f \circ g$ with components:

• •$g_k(\mathbf{x})$ for k = 1..m
• •$f_i(\mathbf{u})$ for i = 1..p
• •$h_i(\mathbf{x}) = f_i(g(\mathbf{x}))$

Then for each output component i and input component j:

This is literally matrix multiplication in coordinates.

Path-tracing interpretation:

• •Each term $\frac{\partial f_i}{\partial u_k}\,\frac{\partial g_k}{\partial x_j}$ corresponds to a path $x_j \to u_k \to y_i$.
• •You add contributions from all intermediate coordinates $u_k$.

“Computational graph” view (interactive canvas)

Think of a small graph:

• •Nodes are variables (scalars or vectors).
• •Edges are functions.

For our two-layer composition:

x →(g)→ u →(f)→ y

On an interactive canvas, you can show two overlays:

Overlay 1: Jacobians on edges

• •edge x→u labeled $Dg(\mathbf{x})$ (m×n)
• •edge u→y labeled $Df(\mathbf{u})$ (p×m)

To get total derivative x→y, multiply along the path:

Overlay 2: live perturbations and gradients

Pick a concrete point x₀.

• •Drag x by a tiny $\Delta \mathbf{x}$. The canvas updates u and y and also shows the predicted linear response $Df\,Dg\,\Delta \mathbf{x}$.
• •Alternatively, set a scalar loss $\ell = \phi(\mathbf{y})$ (or just take p=1). Show a gradient vector at y and animate it flowing backward:

This “two-way animation” is the visualization you want to internalize:

• •Forward: $\Delta$ vectors multiply by Jacobians.
• •Backward: gradients multiply by Jacobian transposes.

Scalar output special case (gradient form)

Now suppose $f: \mathbb{R}^m \to \mathbb{R}$ is scalar. Then $Df(\mathbf{u})$ is 1×m.

From the Jacobian chain rule:

This left side is a 1×n row vector. If we want the gradient as an n×1 column vector, transpose:

\begin{align*}

\nabla (f\circ g)(\mathbf{x})

&= D(f\circ g)(\mathbf{x})^\top \\

&= \bigl(Df(\mathbf{u})\,Dg(\mathbf{x})\bigr)^\top \\

&= Dg(\mathbf{x})^\top\,Df(\mathbf{u})^\top \\

&= Dg(\mathbf{x})^\top\,\nabla f(\mathbf{u}).

\end{align*}

That last line is the standard “gradient chain rule” used everywhere in ML.

One more composition (three layers)

If you have $\mathbf{x} \xrightarrow{g} \mathbf{u} \xrightarrow{f} \mathbf{v} \xrightarrow{r} \mathbf{y}$, then:

Forward-mode perturbations multiply left-to-right in the same order as the functions apply (inner to outer).

Reverse-mode gradients multiply by transposes right-to-left:

This is backprop in one line—just applied repeatedly.

Why this matters for ML

Neural networks are compositions of many vector-valued functions:

Training requires $\nabla_{\theta}\ell$, gradients with respect to parameters. The only tool you need conceptually is the multivariable chain rule, but applied efficiently.

The hard part is not the calculus; it’s bookkeeping:

• •What depends on what?
• •What is the shape of each Jacobian?
• •Are we pushing perturbations forward or pulling gradients back?

A concrete computational graph (with explicit shapes)

Let’s build a tiny “network” with one hidden layer and a scalar loss. Define:

• •x ∈ ℝ²
• •Parameters:
• •$W$ ∈ ℝ^{3×2}, b ∈ ℝ³
• •c ∈ ℝ³ (a vector used to reduce to scalar)

Forward computation:

1) a = $W\mathbf{x} + \mathbf{b}$ (so a ∈ ℝ³)

2) h = $\sigma(\mathbf{a})$ elementwise (so h ∈ ℝ³)

3) $\ell = \mathbf{c}^\top \mathbf{h}$ (so $\ell$ ∈ ℝ)

This is a composition:

x →(affine)→ a →(nonlinearity)→ h →(dot)→ ℓ

On an interactive canvas, you can attach:

Forward-mode (push a perturbation)

A perturbation $\Delta \mathbf{x}$ produces:

\begin{align*}

\Delta \mathbf{a} &\approx W\,\Delta \mathbf{x} \\

\Delta \mathbf{h} &\approx \operatorname{Diag}(\sigma'(\mathbf{a}))\,\Delta \mathbf{a} \\

\Delta \ell &\approx \mathbf{c}^\top\,\Delta \mathbf{h}.

\end{align*}

Combine:

So the total Jacobian (1×2 row vector) is:

Reverse-mode (pull back a gradient)

Because ℓ is scalar, we typically want $\nabla_{\mathbf{x}}\ell$ as a 2×1 column vector.

Start with $\nabla_{\ell}\ell = 1$.

• •From ℓ = cᵀh:

• •Through h = σ(a) elementwise:

• •Through a = $W\mathbf{x}+\mathbf{b}$:

Compare with the forward-mode Jacobian expression above:

• •$D_{\mathbf{x}}\ell$ (row) = cᵀ Diag(σ′) W
• •$\nabla_{\mathbf{x}}\ell$ (column) = $W^\top$ Diag(σ′) c

They are transposes, consistent with $\nabla \ell = (D\ell)^\top$.

Visual intuition: pushforward vs pullback

To address visualization explicitly, here’s the picture you should rehearse:

1) Pick a point x₀.

2) Draw a tiny arrow $\Delta \mathbf{x}$ at x.

3) Multiply by local Jacobians to watch the arrow morph:

• •it rotates/scales/shears in a-space,
• •then again in h-space,
• •finally collapses to a scalar change $\Delta \ell$.

Now reverse:

1) Draw a gradient arrow at h: it points in the direction that increases ℓ fastest in h-space.

2) Pull it back to a using the transpose of the local Jacobian: $\operatorname{Diag}(\sigma')$ (symmetric here, so transpose doesn’t change it).

3) Pull it back to x using $W^\top$.

This is not two unrelated processes. It’s the same linear maps viewed from two dual perspectives:

• •perturbations: $\Delta$ vectors (tangent vectors)
• •gradients: covectors, pulled back by transpose

You don’t need the formal differential-geometry language to use it correctly, but you do need the operational rule:

If forward uses $J$, backward uses $J^\top$.

Connection you’ll use next

Backpropagation is essentially the repeated application of:

where $J$ is the Jacobian of a local block in the computational graph.

Once you’re comfortable multiplying Jacobians (forward) and multiplying by transposes (backward), you’re ready to study backprop as an algorithmic optimization: reuse intermediate results so you don’t form huge Jacobian matrices explicitly.