Matrix Calculus
Gradient of loss wrt weights, numerator/denominator layout
You have linear regression with loss L = ‖Xw - y‖² = (Xw - y)^T(Xw - y). You want ∂L/∂w to take the gradient step. Expanding: L = w^T X^T X w - 2y^T Xw + y^T y. The gradient is ∇_w L = 2X^T Xw - 2X^T y. Setting to zero gives X^T Xw = X^T y — the normal equations. You just derived the closed-form solution to linear regression using matrix calculus. Without it, you would be working element-wise and making sign errors constantly.
Gradient of a scalar by a vector: ∂f/∂x ∈ ℝⁿ, the same shape as x. Key identities: ∂(x^T a)/∂x = a. ∂(x^T A x)/∂x = (A + A^T)x. For symmetric A: 2Ax. These identities cover almost everything in linear models, regularized regression, and Kalman filters.
Jacobian: the derivative of a vector-valued function. If f: ℝⁿ → ℝᵐ, then J = ∂f/∂x ∈ ℝ^{m×n} where J_{ij} = ∂f_i/∂x_j. The Jacobian of the softmax is (diag(s) - ss^T) where s = softmax(x). This is what the backward pass through softmax must compute — a matrix product, not a scalar multiplication.
Chain rule in matrix form: ∂L/∂X = ∂L/∂Y · ∂Y/∂X. The dimensions must work out: if Y = f(X), the Jacobian ∂Y/∂X has shape (dim Y × dim X). For a neural network linear layer z = Wx + b, the gradient with respect to W is ∂L/∂W = (∂L/∂z) · x^T — the outer product of the upstream gradient and the input. This one identity covers every fully-connected layer.
NOT this. You do not need matrix calculus if you use autograd — this is false. Autograd computes gradients correctly, but you need matrix calculus to debug shape errors in custom operations, to verify that backpropagation through a novel layer is correct, and to understand why certain operations are expensive to differentiate. Every custom PyTorch layer that implements a backward() method is matrix calculus. If you cannot derive the gradient manually, you cannot verify that your custom backward pass is correct.
Key points
- Memorize the identity ∇_w (Aw)^T B(Aw) = 2A^T BA w for symmetric B. This is the gradient of a quadratic form and appears in every regularized linear model, least-squares problem, and Kalman filter update. Setting it to zero gives the normal equations. Deriving it from scratch each time is error-prone and slow.
- Trap: layout convention inconsistency. Numerator layout and denominator layout conventions differ between textbooks — the Matrix Cookbook (Petersen & Pedersen) uses denominator layout; most ML papers use numerator layout. Mixing conventions within a derivation produces a Jacobian that is transposed relative to what you need, giving a gradient update applied in the wrong direction. Pick one convention and never mix.
- Diagnostic: if your manually implemented backward pass does not match autograd's gradient, try transposing the Jacobian. Layout convention errors are the most common cause of manual backward pass bugs — the gradient has the right values but the wrong shape. The fix is to check whether you need J or J^T in the chain rule expression, which depends on your chosen layout convention.
- The trace trick: for a scalar built from tr(·), write its differential, use tr(M^T)=tr(M) and cyclic invariance tr(ABC)=tr(BCA) to collect every term into the form tr(G^T dX), then the gradient is ∂f/∂X = G. For f(X) = tr(X^T B X): df = tr((dX)^T BX) + tr(X^T B dX) = tr((BX)^T dX) + tr(X^T B dX) = tr((BX + B^T X)^T dX), so ∂f/∂X = (B+B^T)X — the matrix generalization of the vector identity ∂(x^T A x)/∂x = (A+A^T)x above, and the technique checkQuestion 2 below asks you to apply.
The gradient of a scalar loss with respect to any weight matrix is the outer product of the upstream gradient and the input activation. That one pattern covers every fully-connected layer. Layout convention errors are the most common silent bug in custom backpropagation.
Recap
- Normal equations from matrix calculus: ∇_w‖Xw−y‖² = 2XᵀXw − 2Xᵀy = 0 ⟹ XᵀXw = Xᵀy.
- Key identities: ∂(xᵀa)/∂x = a; ∂(xᵀAx)/∂x = (A+Aᵀ)x; for symmetric A, 2Ax.
- Jacobian J ∈ ℝ^{m×n}, J_{ij}=∂f_i/∂x_j; softmax Jacobian = diag(s) − ssᵀ.
- Linear-layer gradient = outer product: ∂L/∂W = (∂L/∂z)·xᵀ — covers every fully-connected layer.
- Autograd doesn't remove the need — you need matrix calculus to write and debug any custom backward().
- Layout convention is the silent bug: numerator vs denominator layout differ; mixing gives a transposed Jacobian.
- Backward-pass mismatch with autograd? Try transposing the Jacobian — layout errors (right values, wrong shape) are the most common cause.
Check your understanding
Q1. The Jacobian of a function f: ℝⁿ → ℝᵐ at point x can be written in numerator layout or denominator layout. Which two of the following give a self-consistent, correct description of J under one of these conventions?
- A) J is m×n. Element J_{ij} = ∂fᵢ/∂xⱼ — the partial derivative of the i-th output wrt the j-th input. J is the best linear approximation to f near x: f(x+δ) ≈ f(x) + J·δ. For a scalar function (m=1), J reduces to the gradient ∇f, a 1×n row vector. In backpropagation: ∂L/∂x = Jᵀ·(∂L/∂f) — the vector-Jacobian product reverse-mode autodiff computes, which is what loss.backward() returns.
- B) J is n×m. Element J_{ij} = ∂xᵢ/∂fⱼ — how much the j-th output changes the i-th input, measuring the inverse sensitivity, since input-output roles are swapped relative to the standard convention. For a scalar loss (m=1), J is a column vector ∈ ℝⁿ. The backpropagation update uses J directly, without transposing, because the gradient is assumed to flow in the same direction as the Jacobian's column structure.
- C) J is n×m, the transpose of the numerator-layout Jacobian: J_{ij} = ∂fⱼ/∂xᵢ. This is the denominator layout convention — rows are inputs, columns are outputs. For a scalar function (m=1), J is an n×1 column vector identical to ∇f. Under this convention, ∂L/∂x = J·(∂L/∂f) multiplies J directly by the upstream gradient without transposing, since the transpose is already absorbed into J's definition.
- D) J is m×m — a square matrix regardless of input/output dimensions, since it represents covariance between outputs rather than a mapping from inputs. Element J_{ij} = ∂fᵢ/∂fⱼ measures how outputs co-vary. The diagonal elements J_{ii} = 1 always (each output is perfectly correlated with itself). Off-diagonal elements capture how changing one output requires changing another, which determines gradient flow between neurons.
Q2. Compute ∂/∂W(tr(WᵀAW)) where A is symmetric n×n and W is n×k.
- A) ∂/∂W tr(WᵀAW) = AW (not 2AW). This follows from the product rule: d(WᵀAW) = (dW)ᵀAW + WᵀA(dW), so tr(d(WᵀAW)) = tr((dW)ᵀAW) + tr(WᵀA dW) — but treating tr((dW)ᵀAW) as already identical to tr(WᵀA dW) rather than transposing it first (tr((dW)ᵀAW) = tr((AW)ᵀdW) = tr(WᵀAᵀdW)) silently drops one of the two terms as a duplicate, undercounting the gradient by a factor of 2.
- B) ∂/∂W tr(WᵀAW) = WᵀA + AW. Using the identity ∂tr(XᵀBX)/∂X = (B+Bᵀ)X: here X=W, B=A (symmetric), so gradient = 2AW. But the denominator layout convention instead gives Wᵀ·A on the left: (WᵀA)ᵀ = AW, so both layout conventions happen to agree on 2AW in this symmetric case.
- C) Let f(W) = tr(WᵀAW). Then df = tr((dW)ᵀAW) + tr(WᵀA dW). Since tr((dW)ᵀAW) = tr(WᵀAᵀdW) = tr(WᵀAdW) using A=Aᵀ, df = tr(2WᵀA dW) = tr((2AW)ᵀdW), giving ∂f/∂W = 2AW. This appears in PCA: max_W tr(WᵀΣW) s.t. WᵀW=I gives gradient 2ΣW, leading via Lagrange multipliers to the eigenvector equation ΣW=WΛ.
- D) ∂/∂W tr(WᵀAW) = tr(A)·W. The trace of a product tr(WᵀAW) = tr(A)·tr(WᵀW) when A and WᵀW are both symmetric matrices that happen to share the same eigenbasis, so the gradient is tr(A)·∂tr(WᵀW)/∂W = tr(A)·2W, treating the trace of a matrix product as always factoring into the product of the individual traces.
Q3. What is the gradient of the softmax cross-entropy loss with respect to the pre-softmax logits z? Derive the clean form.
- A) ∂L/∂z_k = p_k·(1−p_k) for k=y (true class) and −p_k·p_y for k≠y. This follows from the softmax Jacobian: ∂pᵢ/∂zⱼ = pᵢ(δᵢⱼ−pⱼ). Applying chain rule with ∂L/∂pᵢ = −1/p_y·δᵢy gives the above. The gradient is bounded by the product of probabilities, explaining why softmax prevents gradient saturation for the cross-entropy loss, unlike the unbounded gradients that pure MSE would produce through the same softmax layer.
- B) ∂L/∂z = −e_y/p_y where e_y is a one-hot vector. The gradient of cross-entropy L = −log(p_y) with respect to z is: ∂L/∂z_k = ∂/∂z_k(−log p_y) = −(1/p_y)·∂p_y/∂z_k. For k=y: ∂p_y/∂z_y = p_y(1−p_y), so ∂L/∂z_y = −(1−p_y). For k≠y: ∂p_y/∂z_k = −p_y·p_k, so ∂L/∂z_k = p_k. In vector form: ∂L/∂z = p − e_y, which contradicts the −e_y/p_y form stated at the start of this derivation.
- C) ∂L/∂z_k = p_k for all k, including the true class k=y. The softmax cross-entropy gradient simply equals the predicted probability vector p because the derivative of log(softmax(z)) with respect to z is (I − 11ᵀ/n)·p in the numerator layout convention. Setting the gradient to zero means all probabilities must be equal — the uniform distribution 1/n at every coordinate is the unique critical point of this loss surface.
- D) Let p = softmax(z), pᵢ = e^{zᵢ}/Σⱼe^{zⱼ}. Cross-entropy: L = −log(p_y) = −z_y + log(Σⱼe^{zⱼ}). For k=y: ∂L/∂z_y = −1 + p_y = p_y−1. For k≠y: ∂L/∂z_k = p_k. Combined: ∂L/∂z = p − e_y, the probability vector minus the one-hot label. This clean form arises because the cross-entropy's 1/p derivative cancels the softmax's p(1−p) Jacobian term, eliminating saturation.
Try it interactively
ML Systems Lab is a free interview-prep platform for ML engineers — work through the full interactive module, quizzes, and drills.
Open ML Systems Lab →