Matrix calculus and gradients

Training a machine learning model means adjusting its parameters to reduce a loss function, and the tool for doing that is the gradient: the vector of partial derivatives that points in the direction of steepest increase. A modern neural network has millions or billions of parameters, and every training step computes the gradient of the loss with respect to all of them. The method that makes this feasible, backpropagation, is nothing more than the chain rule of calculus organized efficiently, and it is what deep learning frameworks such as PyTorch and JAX automate.

This chapter builds the calculus needed to understand it: gradients, directional derivatives, Jacobians and Hessians; the handful of matrix derivative identities used constantly; the least-squares solution derived by setting a gradient to zero; the gradient of the softmax cross-entropy loss; and backpropagation through a small network, checked numerically.

Gradients and directional derivatives

For a scalar function of a vector, \(f: \mathbb{R}^n \to \mathbb{R}\), the gradient collects the partial derivatives:

The rate of change of \(f\) in the direction of a unit vector \(\mathbf{u}\) is the directional derivative \(\nabla f^\top\mathbf{u} = \|\nabla f\|\cos\theta\). It is largest when \(\mathbf{u}\) points along the gradient, so the gradient is the direction of steepest ascent and \(-\nabla f\) the direction of steepest descent; its length is the steepest slope. Gradient descent, \(\mathbf{x} \leftarrow \mathbf{x} - \eta\nabla f(\mathbf{x})\), takes small steps downhill. The gradient is also perpendicular to the contour lines (level sets) of \(f\).

Jacobians and Hessians

For a vector-valued function \(\mathbf{f}: \mathbb{R}^n \to \mathbb{R}^m\), the derivative is the \(m \times n\) Jacobian matrix \(\mathbf{J}\), with \(J_{ij} = \partial f_i/\partial x_j\). It is the best linear approximation of \(\mathbf{f}\) near a point: \(\mathbf{f}(\mathbf{x} + \boldsymbol{\delta}) \approx \mathbf{f}(\mathbf{x}) + \mathbf{J}\boldsymbol{\delta}\). Two cases appear constantly in neural networks:

• A linear layer \(\mathbf{f}(\mathbf{x}) = \mathbf{A}\mathbf{x}\) has Jacobian \(\mathbf{J} = \mathbf{A}\).
• An element-wise function such as \(\tanh\) or ReLU has a diagonal Jacobian, \(\operatorname{diag}(\sigma\prime(x_1), \dots, \sigma\prime(x_n))\), because each output depends only on its own input. In code it is never formed as a matrix; it is simply an element-wise multiplication.

The second derivatives of a scalar function form the Hessian, \(H_{ij} = \partial^2 f/\partial x_i\,\partial x_j\), a symmetric \(n \times n\) matrix describing curvature. It gives the quadratic approximation \(f(\mathbf{x} + \boldsymbol{\delta}) \approx f(\mathbf{x}) + \nabla f^\top\boldsymbol{\delta} + \tfrac12\boldsymbol{\delta}^\top\mathbf{H}\boldsymbol{\delta}\). At a point where the gradient is zero, a positive definite Hessian means a local minimum, a negative definite one a local maximum, and an indefinite one (eigenvalues of both signs) a saddle point.

Identities for matrix derivatives

A small set of results, all derivable from the component definitions, covers most of what is needed. For vectors \(\mathbf{x}, \mathbf{a}\) and a matrix \(\mathbf{A}\):

A useful habit is to check shapes: the gradient of a scalar with respect to \(\mathbf{x} \in \mathbb{R}^n\) must itself be in \(\mathbb{R}^n\), so for example \(\mathbf{A}^\top(\mathbf{A}\mathbf{x} - \mathbf{b})\) has the right shape (\(n \times m\) times \(m\)), while \(\mathbf{A}(\mathbf{A}\mathbf{x} - \mathbf{b})\) would not even be defined for a non-square \(\mathbf{A}\).

Least squares from the gradient

Linear regression fits weights \(\mathbf{w}\) to minimize the squared error \(L(\mathbf{w}) = \|\mathbf{X}\mathbf{w} - \mathbf{y}\|^2\) over \(N\) examples. By the last identity, \(\nabla L = 2\mathbf{X}^\top(\mathbf{X}\mathbf{w} - \mathbf{y})\). Setting it to zero gives the normal equations and their solution:

Because \(L\) is a convex quadratic (its Hessian \(2\mathbf{X}^\top\mathbf{X}\) is positive semidefinite), this stationary point is the global minimum. In practice the system is solved with a factorization rather than an explicit inverse.

The program below checks the least-squares gradient numerically and then solves the normal equations. Checking an analytic gradient against finite differences is the standard way to catch derivation and coding errors.

import numpy as np

rng = np.random.default_rng(0)
X = rng.normal(size=(50, 3))
y = X @ np.array([1.5, -2.0, 0.5]) + 0.1 * rng.normal(size=50)
w = np.zeros(3)

def loss(w):
    r = X @ w - y
    return r @ r                      # sum of squared residuals

grad = 2 * X.T @ (X @ w - y)          # analytic gradient

eps = 1e-6                            # central differences, one coordinate at a time
num = np.array([(loss(w + eps * e) - loss(w - eps * e)) / (2 * eps) for e in np.eye(3)])
print("analytic :", grad.round(4))
print("numerical:", num.round(4))
print("agree:", np.allclose(grad, num, rtol=1e-6))
w_star = np.linalg.solve(X.T @ X, X.T @ y)       # set the gradient to zero
print("least-squares weights:", w_star.round(3))

analytic : [-182.2827  129.706   -60.6253]
numerical: [-182.2827  129.706   -60.6253]
agree: True
least-squares weights: [ 1.518 -2.016  0.502]

The recovered weights (1.518, −2.016, 0.502) are close to the true (1.5, −2.0, 0.5) used to generate the data; the small differences come from the noise.

The chain rule

Most functions in machine learning are compositions: a loss of a prediction of a layer of a layer. The chain rule handles them. For scalars, if \(y = f(u)\) and \(u = g(x)\), then \(dy/dx = (dy/du)(du/dx)\). For vectors, Jacobians multiply: if \(\mathbf{y} = \mathbf{f}(\mathbf{u})\) and \(\mathbf{u} = \mathbf{g}(\mathbf{x})\), then \(\mathbf{J}_{\mathbf{y}\mathbf{x}} = \mathbf{J}_{\mathbf{y}\mathbf{u}}\mathbf{J}_{\mathbf{u}\mathbf{x}}\).

Softmax and cross-entropy

Classifiers with \(K\) classes output \(K\) scores (logits) \(\mathbf{z}\), turned into probabilities by the softmax function, and are trained with the cross-entropy loss against a one-hot target \(\mathbf{y}\) (Lesson 9 explains why):

Working through the derivatives, the softmax and the logarithm combine into a remarkably simple gradient:

The gradient is just the predicted probabilities minus the target: it pushes up the logit of the correct class by how far its probability falls short of 1, and pushes down every other logit by its probability. This clean result is one reason softmax and cross-entropy are always used together.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max())           # subtract the max for numerical stability
    return e / e.sum()

z = np.array([2.0, 1.0, 0.1])         # logits for 3 classes
y = np.array([0.0, 1.0, 0.0])         # true class is class 1
p = softmax(z)
L = -np.sum(y * np.log(p))
print("p =", p.round(4), " loss =", round(L, 4))
print("gradient p - y =", (p - y).round(4))

eps = 1e-6
num = np.array([(-np.sum(y * np.log(softmax(z + eps * e))) - -np.sum(y * np.log(softmax(z - eps * e)))) / (2 * eps)
                for e in np.eye(3)])
print("numerical      =", num.round(4))

p = [0.659  0.2424 0.0986]  loss = 1.417
gradient p - y = [ 0.659  -0.7576  0.0986]
numerical      = [ 0.659  -0.7576  0.0986]

Subtracting the maximum logit before exponentiating does not change the softmax (it cancels in the ratio) but prevents overflow when logits are large, a standard numerical precaution.

Backpropagation

Lesson 5 shows how frameworks automate what follows for any program, and the numerical precision it runs in.

A neural network is a chain of layers, and its loss is a composition of many functions. Backpropagation computes the gradient of the loss with respect to every parameter by applying the chain rule from the output backwards. A forward pass computes and stores each layer's output. A backward pass then starts with \(\partial L/\partial\hat{y}\) and, layer by layer, multiplies by each layer's local derivative, producing at each layer both the gradient for that layer's parameters and the gradient to pass to the layer below.

import numpy as np

rng = np.random.default_rng(42)
x = rng.normal(size=4)                 # one input
t = 1.0                                # target
W1 = rng.normal(size=(3, 4)); W2 = rng.normal(size=3)

def forward(W1, W2):
    h = np.tanh(W1 @ x)                # hidden layer
    yhat = W2 @ h                      # output
    return 0.5 * (yhat - t) ** 2, h, yhat

L, h, yhat = forward(W1, W2)
# Backpropagation: apply the chain rule from the loss backwards
dy = yhat - t                          # dL/dyhat
dW2 = dy * h                           # dL/dW2
dh = dy * W2                           # dL/dh
dz = dh * (1 - h ** 2)                 # through tanh: d tanh(z)/dz = 1 - tanh^2
dW1 = np.outer(dz, x)                  # dL/dW1

eps = 1e-6
E = np.zeros_like(W1); E[1, 2] = eps   # check one entry of dW1
num = (forward(W1 + E, W2)[0] - forward(W1 - E, W2)[0]) / (2 * eps)
print(f"loss {L:.4f}")
print(f"backprop dL/dW1[1,2] = {dW1[1, 2]:.6f}")
print(f"numerical            = {num:.6f}")

loss 3.2841
backprop dL/dW1[1,2] = 0.076549
numerical            = 0.076549

Notice the pattern: the gradient for a weight matrix is always an outer product of the gradient arriving from above and the input that the layer received in the forward pass. That is why training needs memory to store activations, which in large models is often a bigger constraint than storing the weights themselves.

Why reverse mode

There are two ways to apply the chain rule through a long composition. Forward mode carries derivatives along with values from input to output, which costs one pass per input variable. Reverse mode, backpropagation, carries them from output to input, which costs one backward pass per output. A training loss has millions of inputs (parameters) and one output (the loss), so reverse mode computes the full gradient in roughly the time of two or three forward passes, while forward mode would need millions. Automatic differentiation libraries implement reverse mode generally: they record the operations of the forward pass and replay them backwards, so practitioners rarely write gradients by hand. Understanding what they compute, and checking gradients numerically when writing custom operations, remains essential.