Matrix Calculus: Gradients in High Dimensions

Beyond Scalar Calculus

In standard calculus, you differentiate a scalar function \left. f:\mathbb{R}\rightarrow\mathbb{R} \right. The derivative is a single number.

In Machine Learning, we deal with functions of many variables (neural networks have millions of parameters) that output vectors, matrices, or scalars. Matrix Calculus provides the notation and rules to compute these derivatives efficiently.

Why It Matters for ML

Every time you call loss.backward() in PyTorch, you are computing matrix derivatives. Understanding matrix calculus lets you derive gradients by hand, debug backpropagation, and design custom layers.

Layout Conventions

There are two conventions for arranging derivatives: Numerator Layout and Denominator Layout. ML typically uses Numerator Layout (also called Jacobian layout).

Numerator Layout

The derivative has the same shape as the numerator. If y \in \mathbb{R}^{m} and x \in \mathbb{R}^{n} then \frac{\partial y}{\partial x} is m \times n

\text{Shape}(\frac{\partial y}{\partial x}) = \text{Shape}(y) \times \text{Shape}(x)^{T}

Denominator Layout

The derivative has the same shape as the denominator transposed. Used in some optimization textbooks.

\text{Shape}(\frac{\partial y}{\partial x}) = \text{Shape}(x)^{T} \times \text{Shape}(y)

Warning: Always check which convention a paper or library uses. Transposing the wrong matrix leads to dimension mismatch bugs in backprop.

Gradients

The gradient of a scalar function \left. f:\mathbb{R}^{n}\rightarrow\mathbb{R} \right. is a vector of partial derivatives. In numerator layout, it is technically a row vector, but we often treat it as a column vector for convenience in update rules.

\nabla f = \begin{bmatrix} \frac{\partial f}{\partial x_{1}} \\ \frac{\partial f}{\partial x_{2}} \\ {\vdots} \\ \frac{\partial f}{\partial x_{n}} \end{bmatrix}∇f=​∂x1​∂f​∂x2​∂f​⋮∂xn​∂f​​​

The gradient points in the direction of steepest ascent of the function.

Key Gradient Identities

• \nabla_{x}(a^{T}x) = a (Linear function)
• \nabla_{x}(x^{T}x) = 2x (Squared norm, recall norms)
• \nabla_{x}(x^{T}Ax) = (A + A^{T})x (Quadratic form)
• \nabla_{x}\mid\mid Ax - b\mid\mid^{2} = 2A^{T}(Ax - b) (Least squares)

Jacobian Matrix

When the output is a vector \left. f:\mathbb{R}^{n}\rightarrow\mathbb{R}^{m} \right. we need an m \times n matrix to capture all partial derivatives. This is the Jacobian.

J = \frac{\partial f}{\partial x} = \begin{bmatrix} \frac{\partial f_{1}}{\partial x_{1}} & \cdots & \frac{\partial f_{1}}{\partial x_{n}} \\ {\vdots} & \ddots & {\vdots} \\ \frac{\partial f_{m}}{\partial x_{1}} & \cdots & \frac{\partial f_{m}}{\partial x_{n}} \end{bmatrix}J=∂x∂f​=​∂x1​∂f1​​⋮∂x1​∂fm​​​⋯⋱⋯​∂xn​∂f1​​⋮∂xn​∂fm​​​​

Linear Transformation View

The Jacobian tells you how the output space locally stretches and rotates. If f(x) = Ax then J = A The derivative of a linear map is the map itself!

Interactive: Jacobian as Linearization

Explore how the Jacobian matrix acts as a local linear approximation for a non-linear transformation. Notice how the grid lines locally look like the Jacobian vectors.

Jacobian as Linearization

The Jacobian matrix J is the "linear approximation" of a non-linear function.

Input Space (x, y)

x: 1.00, y: 1.00

Output Space f(x, y)

u: 0.79, v: 1.45

Jacobian Matrix at Point (1.00, 1.00)

J =

1.00

-0.91

-0.42

1.00

Col 1: How output changes when x moves.

Col 2: How output changes when y moves.

Hessian Matrix

The Hessian is the matrix of second derivatives. It tells you about the curvature of a scalar function.

H = \nabla^{2}f = \begin{bmatrix} \frac{\partial^{2}f}{\partial x_{1}^{2}} & \frac{\partial^{2}f}{\partial x_{1}\partial x_{2}} & \cdots \\ \frac{\partial^{2}f}{\partial x_{2}\partial x_{1}} & \frac{\partial^{2}f}{\partial x_{2}^{2}} & \cdots \\ {\vdots} & {\vdots} & \ddots \end{bmatrix}H=∇2f=​∂x12​∂2f​∂x2​∂x1​∂2f​⋮​∂x1​∂x2​∂2f​∂x22​∂2f​⋮​⋯⋯⋱​​

H positive definite

Local minimum. Bowl curving up in all directions (convex).

H negative definite

Local maximum. Bowl curving down (concave).

H indefinite

Saddle point. Curves up in some directions, down in others.

H singular

Flat direction. Need higher order derivatives.

Interactive Rules

Explore the three fundamental rules of matrix calculus: scalar-by-vector, vector-by-vector (Jacobian), and the chain rule.

Matrix Calculus Rules

Interactive guide to common derivative identities.

Scalar by Vector

Vector by Vector (Jacobian)

Chain Rule (Composition)

∂(wᵀx)/∂x = w

The gradient of a dot product wᵀx with respect to x is simply w.

f(x) = wᵀx

→

∇f = w

Practical Example

If f(x) = 3x₁ + 2x₂, then ∇f = [3, 2]ᵀ

Quick Reference Cheat Sheet

∂(xᵀx)/∂x2x

∂(xᵀAx)/∂x(A + Aᵀ)x

∂(aᵀXb)/∂Xabᵀ

∂log|X|/∂XX⁻ᵀ

∂tr(AX)/∂XAᵀ

∂||Ax-b||²/∂x2Aᵀ(Ax-b)

Matrix Chain Rule

The chain rule in matrix calculus is about multiplying Jacobians in the right order.

\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial x}

For L = f(g(x)) chain the Jacobians. Be careful with dimensions!

Backpropagation is Chain Rule

A neural network is a composition of functions: L = f_{n}(f_{n - 1}(...f_{1}(x)))

Backprop computes gradients by starting from the loss and multiplying Jacobians backwards:

\frac{\partial L}{\partial x} = \frac{\partial L}{\partial f_{n}} \cdot \frac{\partial f_{n}}{\partial f_{n - 1}}\cdots\frac{\partial f_{1}}{\partial x}

Case Study: Linear Regression

The Problem

Predict bulb lifespan (y) from features (voltage, temperature, filament thickness). We have data matrix X \in \mathbb{R}^{n \times d} and targets y \in \mathbb{R}^{n} Find optimal weights w

The Loss Function

L(w) = \mid\mid Xw - y\mid\mid^{2} = (Xw - y)^{T}(Xw - y)

The Gradient

Using matrix calculus rules:

\nabla_{w}L = 2X^{T}(Xw - y)

Setting to zero gives the Normal Equations: X^{T}Xw = X^{T}y This is an exact solution!

ML Applications

Custom Autograd

When implementing a custom layer (e.g., in PyTorch or JAX), you must define the backward pass. This requires analytically deriving the Jacobian of your operation.

Natural Gradient Descent

Uses the Fisher Information Matrix (expected Hessian of log-likelihood) to precondition gradients, correcting for the geometry of the parameter space.

Variational Inference

Optimizing the ELBO in VAEs requires computing derivatives of log-determinants and trace operations, which are pure matrix calculus.

Gaussian Processes

Optimizing kernel hyperparameters requires gradients of the marginal likelihood, involving derivatives of matrix inverses and determinants: \frac{\partial K^{- 1}}{\partial\theta} = - K^{- 1}\frac{\partial K}{\partial\theta}K^{- 1}