Gradients of Vector-Valued Functions

When $f : \mathbb{R}^n \to \mathbb{R}^m$ maps vectors to vectors, the analog of the gradient is a matrix called the Jacobian, $J_f = \frac{\partial f}{\partial \mathbf{x}} \in \mathbb{R}^{m \times n}$. Its $(i,j)$ entry is $\frac{\partial f_i}{\partial x_j}$ — the sensitivity of the $i$-th output component to the $j$-th input. Row $i$ is the gradient of $f_i$ transposed.

The Jacobian is the best linear approximation of $f$ near a point: $f(\mathbf{x} + \Delta \mathbf{x}) \approx f(\mathbf{x}) + J_f(\mathbf{x}) \Delta \mathbf{x}$. This extends the 1D idea $f(x + h) \approx f(x) + f'(x) h$ to arbitrary input/output dimensions. It's why we call $f$ 'locally linear' and why gradient-based methods work — the nonlinear function behaves like a matrix at each point.

The chain rule generalizes elegantly to Jacobians. If $\mathbf{z} = f(\mathbf{y})$ and $\mathbf{y} = g(\mathbf{x})$, then $\frac{\partial \mathbf{z}}{\partial \mathbf{x}} = \frac{\partial \mathbf{z}}{\partial \mathbf{y}} \cdot \frac{\partial \mathbf{y}}{\partial \mathbf{x}} = J_f \cdot J_g$ — just multiply Jacobians in order. This is the mathematical backbone of backpropagation: each layer contributes one Jacobian factor, and the chain rule stitches them together.

For linear functions $f(\mathbf{x}) = A\mathbf{x} + \mathbf{b}$, the Jacobian is simply $A$ — the linear approximation is exact. For element-wise nonlinearities like $\sigma$ applied to a vector, the Jacobian is $\text{diag}(\sigma'(\mathbf{x}_1), \dots, \sigma'(\mathbf{x}_n))$. These two cases, combined via chain rule, give you the Jacobian of any feedforward network.

In ML, the Jacobian drives backpropagation (chaining layer-by-layer Jacobians), adversarial attacks (linearizing a classifier via its input Jacobian), neural ODE solvers (integrating flow Jacobians for invertibility), and normalizing flows (tracking $\log|\det J|$ for density estimation). Working comfortably with Jacobians is the bridge from scalar calculus to the matrix calculus that dominates modern deep learning.

Worked example — Jacobian of $f(x,y) = (xy,\; x+y,\; x^2)$: The map goes $\mathbb{R}^2 \to \mathbb{R}^3$, so $J_f$ is $3\times 2$. Taking partials row by row: $J_f(x,y) = \begin{pmatrix}y & x \\ 1 & 1 \\ 2x & 0\end{pmatrix}$. At $(x,y) = (1,2)$ this is $\begin{pmatrix}2 & 1 \\ 1 & 1 \\ 2 & 0\end{pmatrix}$ — the dimensions check out ($m=3$ rows of outputs, $n=2$ columns of inputs).