Matrix calculus for ML

Summary

Matrix calculus extends the ordinary derivative to functions whose inputs and/or outputs are vectors or matrices. The two organizing objects are the gradient (for scalar-valued functions) and the Jacobian (for vector-valued functions); the Hessian is the matrix of second partial derivatives.

Every learning algorithm computes derivatives of a scalar loss with respect to vector- or tensor-valued parameters. Knowing the right shapes and conventions saves hours of debugging. Backpropagation is matrix calculus applied recursively through a computation graph.

The four shape combinations

$f:{\mathbb{R}}^n\to \mathbb{R}$	$f:{\mathbb{R}}^n\to {\mathbb{R}}^m$
Gradient $\nabla f\in {\mathbb{R}}^n$	Jacobian $J_f\in {\mathbb{R}}^{m\times n}$, $(J_f{)}_{ij}=\partial f_i/\partial x_j$
Second derivative: Hessian ${\nabla }^{2}f\in {\mathbb{R}}^{n\times n}$, $({\nabla }^{2}f{)}_{ij}={\partial }^{2}f/\partial x_i\partial x_j$	(rarely used; tensor-valued)

Two competing conventions exist:

• Numerator layout (used in physics, calculus textbooks): $\nabla f$ is a column vector matching $x$.
• Denominator layout (used in stats, ML): same; modern ML universally uses gradient = column vector with the same shape as the parameter.

For matrix parameters $W\in {\mathbb{R}}^{m\times n}$, the gradient $\partial L/\partial W\in {\mathbb{R}}^{m\times n}$ has the same shape as $W$.

Identities you actually use

Function	Gradient
$f(x)=a^{\top }x$	$\nabla f=a$
$f(x)=x^{\top }Ax$, $A$ symmetric	$\nabla f=2Ax$
$f(x)=\Vert x{\Vert }_2^2$	$\nabla f=2x$
$f(W)=a^{\top }Wb$	$\partial f/\partial W=a{b}^{\top }$
$f(W)=\Vert W{\Vert }_F^2=\mathrm{tr}(W^{\top }W)$	$\partial f/\partial W=2W$
$f(W)=\log \det W$ ($W$ PD)	$\partial f/\partial W=W^{-\top }$
$f(W)=\mathrm{tr}(AW)$	$\partial f/\partial W=A^{\top }$

Chain rule

For $L(\theta )=g(f(\theta ))$:

• Scalar $\to$ scalar $\to$ scalar: $\frac{dL}{d\theta }=g^{\prime }(f(\theta ))\cdot f^{\prime }(\theta )$.
• Vector $\to$ vector $\to$ scalar: ${\nabla }_{\theta }L=J_f^{\top }\cdot {\nabla }_{y}g$ where $y=f(\theta )$.

Backprop is just this chain rule applied to a computation graph layer-by-layer, with the Jacobian-vector product computed implicitly to avoid materializing the full Jacobian.

Learning objective

Use shapes to see how the transposed Jacobian pulls two output gradients back to three parameters.

Hessian and second-order methods

The Hessian ${\nabla }^{2}L$ describes local curvature. Second-order methods (Newton’s method, K-FAC, Shampoo) use it to scale gradients by inverse curvature: $\theta \leftarrow \theta -({\nabla }^{2}L{)}^{-1}\nabla L$.

In modern deep learning, the Hessian is too large to materialize ($O(P^2)$ for $P$ parameters). Approximations:

• Diagonal: keep only diagonal entries (RMSProp, Adam’s $v_t$ approximates this).
• Block-diagonal: per-layer Fisher information (K-FAC).
• Hessian-vector products: ${\nabla }^{2}L\cdot v$ via two backward passes. Used in conjugate gradient, influence functions.

Common pitfalls

• Mismatched layout conventions. Always check whether your reference uses numerator or denominator layout; the difference is a transpose.
• Treating gradients as having the same shape as the loss. They have the shape of the parameter, not the loss.
• Computing Jacobians explicitly. For a vector-to-vector function with $m,n$ both large, the Jacobian is $mn$ entries. Use vector-Jacobian products via torch.autograd.grad or jax.vjp instead.
• Forgetting bias dimensions. $Wx+b$ has $\partial L/\partial b={\sum }_i\partial L/\partial y_i$ summed over the batch dimension.