Gradients of Matrices

In ML, loss functions often depend on whole matrices of parameters (think weight matrices in a neural net layer). The gradient with respect to a matrix is just another matrix: $\frac{\partial L}{\partial W} \in \mathbb{R}^{m \times n}$ has the same shape as $W$, with entry $(i,j)$ giving $\frac{\partial L}{\partial w_{ij}}$.

Matrix calculus identities streamline derivations. The most-used ones: $\nabla_{\mathbf{x}}(\mathbf{a}^\top \mathbf{x}) = \mathbf{a}$, $\nabla_{\mathbf{x}}(\mathbf{x}^\top A \mathbf{x}) = (A + A^\top)\mathbf{x}$ (and $= 2A\mathbf{x}$ when $A$ is symmetric), $\nabla_W(\|X - WY\|_F^2) = -2(X - WY)Y^\top$, and $\nabla_W(\text{tr}(W^\top A)) = A$. These let us differentiate common ML losses quickly without expanding index-by-index.

Using these identities, we can rederive the normal equations. For $L(\mathbf{w}) = \|\mathbf{y} - X\mathbf{w}\|^2$, expanding gives $L = \mathbf{y}^\top \mathbf{y} - 2\mathbf{w}^\top X^\top \mathbf{y} + \mathbf{w}^\top X^\top X \mathbf{w}$. Applying the identities: $\nabla_{\mathbf{w}} L = -2X^\top \mathbf{y} + 2X^\top X \mathbf{w}$. Setting this to zero yields $X^\top X \mathbf{w} = X^\top \mathbf{y}$ — the normal equations, arrived at without a single index calculation.

The denominator layout (also called 'gradient layout') matches the shape of $W$ — the convention most ML texts and auto-diff libraries use. An alternative numerator layout transposes. Be consistent, pick one, and note PyTorch/JAX use denominator layout by default (the gradient of a loss has the same shape as the parameter).

Neural network training is one long matrix gradient computation. A single fully connected layer $\mathbf{h} = \sigma(W\mathbf{x} + \mathbf{b})$ has $\nabla_W L = \delta \mathbf{x}^\top$ where $\delta$ is the upstream gradient after passing through $\sigma'$. This outer-product structure is general: the gradient of a matrix that appears in a linear-plus-nonlinear layer is always an outer product of the downstream error and the upstream activation. This is the atom of deep-learning gradient math.

Worked example — verify $\nabla_{\mathbf{x}}(\mathbf{x}^\top A\mathbf{x}) = (A + A^\top)\mathbf{x}$: Take $A = \begin{pmatrix}1 & 2 \\ 0 & 3\end{pmatrix}$ and $\mathbf{x} = (1, 1)^\top$. Expand $\mathbf{x}^\top A \mathbf{x} = x_1^2 + 2x_1 x_2 + 3x_2^2$, giving gradient $(2x_1 + 2x_2,\; 2x_1 + 6x_2) = (4, 8)$ at $\mathbf{x}=(1,1)$. Alternatively, $A + A^\top = \begin{pmatrix}2 & 2 \\ 2 & 6\end{pmatrix}$, so $(A+A^\top)\mathbf{x} = (4,8)$ — the identity checks out.