Jacobians and Automatic Differentiation

A network maps a vector of parameters to a scalar loss. Writing out every derivative would be wasteful. Automatic differentiation applies local derivative rules to the executed computation and propagates only the products needed for that loss.

One derivative for every input-output pair

For $F:\mathbb R^n\to\mathbb R^m$, the Jacobian has shape $m\times n$:

$J_F(x)_{ij}=\frac{\partial F_i}{\partial x_j},\qquad F(x+h)\approx F(x)+J_F(x)h.$

Rows correspond to outputs and columns to inputs. If $F(x,y)=(x^2+y,xy)$, then

$J_F(x,y)=\begin{bmatrix}2x&1\\y&x\end{bmatrix}.$

At $(1,2)$, a perturbation $(0.01,0)$ predicts output changes $(0.02,0.02)$. The first output's exact change has an additional $0.0001$ from the square.

Composition becomes multiplication

For $F:\mathbb R^n\to\mathbb R^m$ and $G:\mathbb R^m\to\mathbb R^p$,

$J_{G\circ F}(x)=J_G(F(x))J_F(x).$

The shapes are $(p\times m)(m\times n)=p\times n$. The order follows the forward composition. For a scalar loss $L$ after $F$, column gradients travel backward as $\nabla_x L=J_F^T\nabla_{F}L$.

A square Jacobian's determinant measures signed local area or volume scaling. A zero determinant means the linear approximation loses a dimension. It does not by itself prove global non-invertibility: $x\mapsto x^3$ is one-to-one even though its derivative vanishes at zero.

Forward mode: carry a tangent

A dual number $a+b\epsilon$ has $\epsilon^2=0$. Multiplication gives $(a+b\epsilon)(c+d\epsilon)=ac+(ad+bc)\epsilon$. The coefficient of $\epsilon$ automatically follows the product rule.

Starting with $x+1\epsilon$, evaluate a program using these rules. The result is $f(x)+f'(x)\epsilon$. In several dimensions, seed a direction $v$ to compute $Jv$, a Jacobian-vector product. One pass computes one input direction; obtaining all columns generally takes $n$ seeds.

Reverse mode: carry an adjoint

Take $a=x^2$, $b=3x$, $L=a+b$. A forward pass records $x,a,b,L$. Start backward with $\bar L=1$. Addition sends 1 to both $a$ and $b$. Their contributions to $x$ are $2x$ and 3, so $\bar x=2x+3$.

A shared variable must accumulate contributions. Replacing one with another loses a path. Reverse mode computes $J^Tv$ for an output seed $v$. For one scalar loss, a single seed gives all parameter derivatives, at a cost comparable to a small number of forward evaluations. Storing or recomputing intermediate values is the main memory tradeoff.

Neither mode chooses a finite-difference step. It differentiates the operations on the executed path, subject to floating-point arithmetic and the derivative conventions of those operations. Discrete decisions and nondifferentiable points still need attention.

Read beyond

Read the equation in context

The survey distinguishes automatic differentiation from symbolic algebra and finite differences. Reverse accumulation reuses intermediate values and sums contributions from all consumers of a node. The displayed rule restates that process in graph notation.

Your turn

Compare a nonlinear image of a small square with the parallelogram predicted by its Jacobian. Shrink the square, then locate three points where the pleat map has zero determinant.

Coding problems