Jacobian
We defined the differential of a scalar function as a
linear map on displacements.
We now extend this to a vector-valued function \(\mathbf{f} : \mathbb{R}^n \to \mathbb{R}^m\). Its
differential is a linear map, and that map is encoded by a matrix, the Jacobian.
Definition: Differentiability (vector-valued)
Let \(U \subseteq \mathbb{R}^n\) be an open set. A function \(\mathbf{f} : U \to \mathbb{R}^m\)
is differentiable at \(\mathbf{x} \in U\) if there exists a linear map
\(L : \mathbb{R}^n \to \mathbb{R}^m\) such that
\[
\mathbf{f}(\mathbf{x} + \mathbf{h}) - \mathbf{f}(\mathbf{x})
= L(\mathbf{h}) + o(\|\mathbf{h}\|) \quad \text{as } \mathbf{h} \to \mathbf{0}.
\]
When such \(L\) exists it is unique (two linear maps differing by \(o(\|\mathbf{h}\|)\) must
coincide, by the same argument as in the
scalar case) and is called the
differential of \(\mathbf{f}\) at \(\mathbf{x}\), denoted
\(d\mathbf{f}_{\mathbf{x}}\) or simply \(\mathbf{f}'(\mathbf{x})\).
Remark. Differentiability is a stronger condition than the existence of all partial
derivatives \(\partial f_i/\partial x_j\). A function can have all partials at a point yet fail
to be differentiable there (the standard counterexample is \(f(x,y) = xy/(x^2+y^2)\) extended by
\(f(0,0) = 0\)). Differentiability implies the existence of all partials. The converse holds
under continuity of the partials in a neighborhood, a fact we state without proof.
Definition: Jacobian
Let \(U \subseteq \mathbb{R}^n\) be an open set and let \(\mathbf{f} : U \to \mathbb{R}^m\) be
differentiable at \(\mathbf{x} \in U\). The Jacobian of \(\mathbf{f}\) at
\(\mathbf{x}\), denoted \(\mathbf{f}'(\mathbf{x})\) or \(J_{\mathbf{f}}(\mathbf{x})\), is the
\(m \times n\) matrix representing the differential \(d\mathbf{f}_{\mathbf{x}}\) in the standard
bases. Its entries are the partial derivatives
\[
J_{ij} = \frac{\partial f_i}{\partial x_j}, \quad 1 \leq i \leq m, \quad 1 \leq j \leq n,
\]
and the defining linear approximation takes the form
\[
\mathbf{f}(\mathbf{x} + d\mathbf{x}) - \mathbf{f}(\mathbf{x})
= \mathbf{f}'(\mathbf{x})\, d\mathbf{x} + o(\|d\mathbf{x}\|) \quad \text{as } d\mathbf{x} \to \mathbf{0}.
\]
(The identification \(J_{ij} = \partial f_i/\partial x_j\) follows from evaluating the linear
approximation along each coordinate direction \(\mathbf{h} = t\,\mathbf{e}_j\), \(t \to 0\).)
Two notations for the derivative-as-operator. The Jacobian
\(\mathbf{f}'(\mathbf{x})\) is simultaneously (i) an \(m \times n\) matrix of partial derivatives
and (ii) a linear map \(\mathbb{R}^n \to \mathbb{R}^m\). Both views are standard, and the
corresponding notations are:
\[
d\mathbf{f} = \mathbf{f}'(\mathbf{x})\, d\mathbf{x}
\quad \text{(matrix-vector product)}
\]
\[
d\mathbf{f} = \mathbf{f}'(\mathbf{x})[d\mathbf{x}]
\quad \text{(linear map applied to an argument)}
\]
For finite-dimensional calculations the two are interchangeable. The bracket form
\(\mathbf{f}'(\mathbf{x})[d\mathbf{x}]\) is the notation of choice once we leave finite dimensions:
for Fréchet derivatives on Banach spaces, for directional derivatives on
manifolds, and for the operator-valued maps \(\mathrm{Ad}_g[X]\), \(\mathrm{ad}_X[Y]\) that appear
in Lie theory and representation theory. This page uses the matrix-product form by default and the
bracket form where operator composition is cleaner (notably in the Chain Rule).
The \(i\)-th row of \(J := J_{\mathbf{f}}(\mathbf{x})\) records the partial derivatives of the
\(i\)-th output component. The \(j\)-th column records the effect of perturbing the \(j\)-th
input coordinate. Writing out
\(d\mathbf{f} = J\, d\mathbf{x}\) component by component for \(n = m = 2\):
\[
\begin{align*}
d\mathbf{f}
&= \begin{bmatrix}
\partial f_1/\partial x_1 & \partial f_1/\partial x_2 \\\\
\partial f_2/\partial x_1 & \partial f_2/\partial x_2
\end{bmatrix}
\begin{bmatrix} dx_1 \\\\ dx_2 \end{bmatrix} \\\\
&= \begin{bmatrix}
(\partial f_1/\partial x_1)\,dx_1 + (\partial f_1/\partial x_2)\,dx_2 \\\\
(\partial f_2/\partial x_1)\,dx_1 + (\partial f_2/\partial x_2)\,dx_2
\end{bmatrix}.
\end{align*}
\]
The Jacobian acts as a linear operator that maps input displacements \(d\mathbf{x}\) to the
corresponding linear part of the output change \(d\mathbf{f}\).
Convention (shape). For \(\mathbf{f} : \mathbb{R}^n \to \mathbb{R}^m\) the Jacobian
\(\mathbf{f}'(\mathbf{x})\) is \(m \times n\). Two special cases recover earlier objects:
- \(m = 1\) (scalar output, so we write \(f\) for \(\mathbf{f}\)): \(f'(\mathbf{x})\) is a
\(1 \times n\) row vector, equal to \((\nabla f)^\top\) in the
gradient convention
established previously.
- \(n = 1\) (curve): \(\mathbf{f}'(\mathbf{x})\) is an \(m \times 1\) column vector,
the velocity vector of the curve.
As a first example, consider the affine map \(\mathbf{f}(\mathbf{x}) = A\mathbf{x}\) with
\(A \in \mathbb{R}^{m \times n}\) constant. Expanding the increment,
\[
\begin{align*}
\Delta \mathbf{f} &= \mathbf{f}(\mathbf{x} + d\mathbf{x}) - \mathbf{f}(\mathbf{x}) \\\\
&= A(\mathbf{x} + d\mathbf{x}) - A\mathbf{x} \\\\
&= A\, d\mathbf{x},
\end{align*}
\]
the linear part is exact and there is no \(o(\|d\mathbf{x}\|)\) remainder. Reading off the
matrix-product form identifies the Jacobian as \(\mathbf{f}'(\mathbf{x}) = A\). The map \(\mathbf{f}\) is
its own linearization, as expected for an affine map.
Local Transformation of Hypervolumes
When \(m = n\), the Jacobian describes how an infinitesimal hypercube in the
input space is mapped to a parallelotope in the output space.
The absolute value of the determinant
of the Jacobian, \(|\det J|\), is the local volume expansion factor, and the
sign of \(\det J\) records whether orientation is preserved. In generative
modeling, specifically Normalizing Flows, this determinant is what allows the
change-of-variables formula to preserve total probability mass under complex nonlinear
invertible transformations of a base distribution.
Chain Rule
In deep learning, a neural network
is modeled as a composition of differentiable layers. The chain rule reduces differentiation of such a
composition to composing the differentials of the individual layers. Because differentials
are linear maps, composition becomes matrix multiplication of Jacobians.
Theorem: Chain Rule
Let \(U \subseteq \mathbb{R}^n\) and \(V \subseteq \mathbb{R}^p\) be open sets, let
\(\mathbf{h} : U \to V\) be differentiable at \(\mathbf{x} \in U\), and let
\(\mathbf{g} : V \to \mathbb{R}^m\) be differentiable at
\(\mathbf{u} := \mathbf{h}(\mathbf{x})\). Then the composition
\(\mathbf{f} = \mathbf{g} \circ \mathbf{h}\) is differentiable at \(\mathbf{x}\), and its
differential is the composition of the individual differentials:
\[
\mathbf{f}'(\mathbf{x})[d\mathbf{x}]
= \mathbf{g}'(\mathbf{h}(\mathbf{x}))\bigl[\,\mathbf{h}'(\mathbf{x})[d\mathbf{x}]\,\bigr].
\]
Equivalently, at the level of Jacobian matrices,
\[
\mathbf{f}'(\mathbf{x}) = \mathbf{g}'(\mathbf{h}(\mathbf{x}))\,\mathbf{h}'(\mathbf{x})
\quad \text{(an } m \times n \text{ matrix, from } m \times p \text{ times } p \times n\text{).}
\]
Proof.
Let \(J_{\mathbf{h}} := \mathbf{h}'(\mathbf{x})\) and
\(J_{\mathbf{g}} := \mathbf{g}'(\mathbf{u})\). We use throughout the fact that in finite
dimensions every linear map \(L\) is automatically bounded. There exists \(C \geq 0\) with
\(\|L\mathbf{v}\| \leq C\|\mathbf{v}\|\) for all \(\mathbf{v}\) (with Euclidean norms one may take
\(C = \sum_{i,j} |L_{ij}|\), since each entry of \(L\mathbf{v}\) is bounded by the absolute sum of
its row times \(\|\mathbf{v}\|\)). The least admissible \(C\) is the operator norm
\(\|L\|_{\mathrm{op}}\). This is what lets us push \(o(\|\cdot\|)\) and
\(O(\|\cdot\|)\) bounds through \(J_{\mathbf{h}}\) and \(J_{\mathbf{g}}\) below.
By differentiability of \(\mathbf{h}\) at \(\mathbf{x}\),
\[
\begin{align*}
\mathbf{h}(\mathbf{x} + d\mathbf{x}) - \mathbf{h}(\mathbf{x})
&= J_{\mathbf{h}}\, d\mathbf{x} + r_{\mathbf{h}}(d\mathbf{x}), \\\\
r_{\mathbf{h}}(d\mathbf{x}) &= o(\|d\mathbf{x}\|).
\end{align*}
\]
Let \(\Delta\mathbf{u} := J_{\mathbf{h}}\, d\mathbf{x} + r_{\mathbf{h}}(d\mathbf{x})\). Note that
\(\|\Delta\mathbf{u}\| = O(\|d\mathbf{x}\|)\). By differentiability of \(\mathbf{g}\) at
\(\mathbf{u}\) (for \(d\mathbf{x}\) small enough, \(\mathbf{u} + \Delta\mathbf{u} \in V\), since
\(V\) is open and \(\|\Delta\mathbf{u}\| = O(\|d\mathbf{x}\|)\)),
\[
\begin{align*}
\mathbf{g}(\mathbf{u} + \Delta\mathbf{u}) - \mathbf{g}(\mathbf{u})
&= J_{\mathbf{g}}\, \Delta\mathbf{u} + r_{\mathbf{g}}(\Delta\mathbf{u}), \\\\
r_{\mathbf{g}}(\Delta\mathbf{u}) &= o(\|\Delta\mathbf{u}\|).
\end{align*}
\]
Composing these and expanding,
\[
\begin{align*}
\mathbf{f}(\mathbf{x} + d\mathbf{x}) - \mathbf{f}(\mathbf{x})
&= J_{\mathbf{g}} J_{\mathbf{h}}\, d\mathbf{x} \\\\
&\quad + \underbrace{J_{\mathbf{g}}\, r_{\mathbf{h}}(d\mathbf{x})}_{o(\|d\mathbf{x}\|)} \\\\
&\quad + \underbrace{r_{\mathbf{g}}(\Delta\mathbf{u})}_{o(\|\Delta\mathbf{u}\|) = o(\|d\mathbf{x}\|)}.
\end{align*}
\]
The first error term is \(o(\|d\mathbf{x}\|)\) because
\(\|J_{\mathbf{g}}\,r_{\mathbf{h}}(d\mathbf{x})\| \leq \|J_{\mathbf{g}}\|_{\mathrm{op}}\,\|r_{\mathbf{h}}(d\mathbf{x})\|
= \|J_{\mathbf{g}}\|_{\mathrm{op}} \cdot o(\|d\mathbf{x}\|) = o(\|d\mathbf{x}\|)\).
The second requires combining \(r_{\mathbf{g}}(\Delta\mathbf{u}) = o(\|\Delta\mathbf{u}\|)\) with
\(\|\Delta\mathbf{u}\| = O(\|d\mathbf{x}\|)\). Since \(d\mathbf{x} \to 0\) forces
\(\Delta\mathbf{u} \to 0\),
\[
\begin{align*}
\frac{\|r_{\mathbf{g}}(\Delta\mathbf{u})\|}{\|d\mathbf{x}\|}
&= \underbrace{\frac{\|r_{\mathbf{g}}(\Delta\mathbf{u})\|}{\|\Delta\mathbf{u}\|}}_{\to\, 0}
\cdot
\underbrace{\frac{\|\Delta\mathbf{u}\|}{\|d\mathbf{x}\|}}_{\text{bounded}} \\\\
&\longrightarrow 0,
\end{align*}
\]
so \(r_{\mathbf{g}}(\Delta\mathbf{u}) = o(\|d\mathbf{x}\|)\) as well (the case
\(\Delta\mathbf{u} = \mathbf{0}\) gives \(r_{\mathbf{g}} = \mathbf{0}\) directly). Identifying the
linear part gives \(\mathbf{f}'(\mathbf{x}) = J_{\mathbf{g}} J_{\mathbf{h}}\).
Order matters: non-commutativity of matrix multiplication
The Jacobian of a composition is a product of Jacobians in a specific order: outer-first,
inner-last. The dimensional count confirms this. If \(\mathbf{x} \in \mathbb{R}^n\),
\(\mathbf{h}(\mathbf{x}) \in \mathbb{R}^p\), and
\(\mathbf{g}(\mathbf{h}(\mathbf{x})) \in \mathbb{R}^m\), then
\[
\underbrace{\mathbf{g}'(\mathbf{h}(\mathbf{x}))}_{m \times p}
\cdot
\underbrace{\mathbf{h}'(\mathbf{x})}_{p \times n}
\]
produces a well-defined \(m \times n\) matrix. The reversed product
\(\mathbf{h}'(\mathbf{x})\,\mathbf{g}'(\mathbf{h}(\mathbf{x}))\) is defined only when \(n = m\),
and even then it is a \(p \times p\) matrix that in general differs from
\(\mathbf{f}'(\mathbf{x})\). This non-commutativity is the reason chain-rule applications in
machine learning must keep careful track of which direction the composition runs.
Backpropagation
We now specialize the chain rule to the situation that drives neural-network training: a
scalar loss computed by a deep composition. This specialization is what the
reverse-mode automatic differentiation
algorithm, commonly called backpropagation in
neural network training,
computes efficiently.
Consider three composed layers, each differentiable, with compatible dimensions:
\[
\mathbf{x} \xrightarrow{\mathbf{f}_1} \mathbb{R}^{p}
\xrightarrow{\mathbf{f}_2} \mathbb{R}^{q}
\xrightarrow{f_3} \mathbb{R},
\]
with \(\mathbf{x} \in \mathbb{R}^n\) playing the role of the parameters and \(f_3\) scalar-valued so
that the final output is the loss
\[
L(\mathbf{x}) = f_3\bigl(\mathbf{f}_2(\mathbf{f}_1(\mathbf{x}))\bigr).
\]
Since the output is scalar, the Jacobian of \(L\) is a \(1 \times n\) row vector. By the
chain rule, it factors as the product
of the layer Jacobians, each evaluated at the corresponding intermediate value:
\[
L'(\mathbf{x}) = f_3'\,\mathbf{f}_2'\,\mathbf{f}_1'
\quad \text{with shapes } (1 \times q)(q \times p)(p \times n).
\]
By the gradient convention, the column
vector \(\nabla L\) is the transpose of this row: \(L'(\mathbf{x}) = (\nabla L)^\top\).
Associativity and evaluation order
Matrix multiplication is associative but not commutative, so the value of the product
\(f_3'\,\mathbf{f}_2'\,\mathbf{f}_1'\) does not depend on the order of the parenthesization, but the
cost of computing it does. The two natural parenthesizations are:
- Reverse mode (backpropagation): left-to-right evaluation
\(\bigl((f_3' \, \mathbf{f}_2')\, \mathbf{f}_1'\bigr)\). In terms of the computation graph, this
corresponds to starting at the output layer and propagating gradients backward
through the network, hence the name.
- Forward mode: right-to-left evaluation
\(\bigl(f_3'\,(\mathbf{f}_2'\, \mathbf{f}_1')\bigr)\), propagating perturbations forward from
inputs toward outputs.
Tracking the shapes through reverse mode:
\[
\begin{align*}
\bigl(f_3'\,\mathbf{f}_2'\bigr)\,\mathbf{f}_1'
&= \bigl[(1 \times q)(q \times p)\bigr](p \times n) \\\\
&= (1 \times p)(p \times n) \\\\
&= (1 \times n) = (\nabla L)^\top.
\end{align*}
\]
Every intermediate result is a row vector of length at most \(\max(p, q, n)\). No product matrix of size
\(p \times n\) or \(q \times n\) is ever formed. Each layer contributes a single
vector-Jacobian product (VJP). A VJP is the product \(\mathbf{v}^\top J\) of a row vector
\(\mathbf{v}^\top\) with a layer's Jacobian \(J\), the row vector standing on the left.
Forward mode, by contrast, accumulates full intermediate Jacobians:
\[
\begin{align*}
f_3'\,\bigl(\mathbf{f}_2'\,\mathbf{f}_1'\bigr)
&= (1 \times q)\bigl[(q \times p)(p \times n)\bigr] \\\\
&= (1 \times q)(q \times n) \\\\
&= (1 \times n) = (\nabla L)^\top.
\end{align*}
\]
The intermediate matrix \(\mathbf{f}_2'\,\mathbf{f}_1'\) has \(q \times n\) entries. The largest
intermediate product reverse mode forms is the final row vector \((\nabla L)^\top\), with \(n\)
entries when \(n\) is the largest dimension (the typical parameter regime). In this matrix count
the saving of reverse mode is therefore a factor of roughly \(q\). In practice neither mode forms
the layer Jacobians. Forward mode pushes one parameter direction \(\mathbf{e}_j\) through the
network per pass and obtains one entry \(f_3'\,\mathbf{f}_2'\,\mathbf{f}_1'\,\mathbf{e}_j\) of the
gradient. Reverse mode pulls a single row vector back through the network once. Each pass costs a
small constant multiple of one evaluation of \(L\), so forward mode needs \(n\) passes to assemble
the full gradient and reverse mode needs one. This is the fundamental reason deep networks are
trained with reverse-mode autodiff rather than forward-mode. For a scalar loss, the saving grows
with the number of parameters.
The VJP and Memory
Materializing a full Jacobian for a layer with, say, \(10{,}000\) inputs and \(10{,}000\)
outputs would require \(10^8\) entries. To avoid this, modern frameworks such as
PyTorch (through its autograd engine) and JAX (through its
\(\mathrm{vjp}\) transformation) implement each layer's differentiation as a
vector-Jacobian product
\[
\mathbf{v}^\top J,
\]
where \(\mathbf{v}^\top\) is the transpose of the gradient flowing back from the downstream
layers. This allows
the system to compute every gradient needed for training without ever materializing a full layer
Jacobian. Avoiding that materialization is a large part of what makes training of "deep" models
feasible.