Matrix Calculus: Gradients of Layers and Losses

A linear layer can contain millions of weights, yet its backward pass fits on a few lines. Matrix calculus lets us derive those lines with their shapes visible, then compare them with numerical derivatives.

Fix a convention first

For scalar $L(x)$ and column $x$, we write $dL=(\nabla_xL)^Tdx$. The gradient has the same shape as $x$. The Jacobian $\partial y/\partial x$ has output coordinates as rows; some texts call this numerator layout. A row derivative of a scalar is the transpose of our column gradient. Read the convention before copying a formula.

For a matrix $W$, the gradient has the shape of $W$ and

$dL=\operatorname{tr}((\nabla_WL)^T dW).$

This trace is just the sum of the entrywise products of the gradient and the perturbation.

Quadratic forms and least squares

For $L(x)=x^TAx$, differentiating each occurrence of $x$ gives

$\nabla_xL=(A+A^T)x.$

Only when $A$ is symmetric may we simplify this to $2Ax$. The antisymmetric part contributes zero to the scalar quadratic form.

Let $r=Xw-y$ and $L=\tfrac12\|r\|^2$. Then $dL=r^Tdr=r^TX\,dw$, so $\nabla_wL=X^T(Xw-y)$. Averaging over $n$ examples adds a factor $1/n$. Omitting the one-half adds a factor 2. A gradient check often catches a lost reduction factor.

Backward through an affine layer

For $z=Wx+b$ and upstream column gradient $g=\nabla_zL$,

$\nabla_xL=W^Tg,\qquad \nabla_WL=gx^T,\qquad \nabla_bL=g.$

If $W$ is $m\times n$, then $g$ is $m\times1$ and $x^T$ is $1\times n$. Their outer product has exactly the weight shape. A batch sums or averages these outer products according to the definition of its loss.

Softmax and cross-entropy

For logits $z$, set $p_i=e^{z_i}/\sum_j e^{z_j}$. The quotient rule gives

$\frac{\partial p_i}{\partial z_j}=p_i(\delta_{ij}-p_j),\qquad J=\operatorname{diag}(p)-pp^T.$

Adding a constant to every logit leaves softmax unchanged. Correspondingly, each row of $J$ sums to zero. In code, subtract the largest logit before exponentiating to avoid overflow.

For targets $y_i\ge0$ with $\sum_i y_i=1$,

$L=-\sum_i y_i\log p_i,\qquad \nabla_zL=p-y.$

Indeed, the $j$th component is $-\sum_i y_i(\delta_{ij}-p_j)=-y_j+p_j$. This holds for both one-hot and soft target distributions. If targets are not normalised, the general expression is $p\sum_i y_i-y$.

Verify shapes and values

Perturb one parameter at a time and compare its central difference with the corresponding gradient entry. Keep the data and any random masks fixed. Avoid activation kinks and try several step sizes. For a large model, a directional check compares $[L(\theta+hv)-L(\theta-hv)]/(2h)$ with $g^Tv$, requiring just two evaluations per direction.

Read beyond

Read the equation in context

Parr and Howard build neural-network derivatives from scalar rules and Jacobians. The affine map shown here is a compact application of those rules. Expand one output in components, then recover the matrix expression and check its dimensions.

Your turn

Select a candidate gradient, then compare every component with central differences. A small error supports the formula at the tested point; try more than one point before trusting a general rule.

{
  "a": [
    1.5,
    -2
  ],
  "x": [
    2,
    -1,
    0.5
  ],
  "shape": [
    2,
    3
  ]
}
Parameter (2 × 3, row order): [2, 1.1, 1.6, 0.3, -0.6, -1.7]

Coding problems