Matrix Calculus: Deriving the OLS Normal Equations and Backprop Through a Linear Layer by Hand

You cannot derive the OLS normal equations, the backpropagation update for a linear layer, or the gradient of the attention score without matrix calculus. Most practitioners learn it piecemeal — a formula here, a trick there. This post builds the full system: scalar-by-vector, vector-by-vector (Jacobian), and scalar-by-matrix, with four complete worked derivations.

Matrix calculus extends scalar calculus to functions involving vectors and matrices. It is the language in which ML derivations are written, and the gap between "I know what a gradient is" and "I can derive it for any layer" is almost entirely matrix calculus.

Notation conventions

We use numerator layout (also called Jacobian layout): if y ∈ ℝᵐ and x ∈ ℝⁿ, then ∂y/∂x ∈ ℝᵐˣⁿ — the derivative has the same shape as the output in the numerator, stacked over the input in the denominator. For a scalar f and a vector x ∈ ℝⁿ: ∂f/∂x ∈ ℝ¹ˣⁿ (a row vector, but often treated as a column vector ∇f for gradient descent). For a scalar f and a matrix X ∈ ℝᵐˣⁿ: ∂f/∂X ∈ ℝᵐˣⁿ — the gradient has the same shape as the matrix.

Scalar by vector: the gradient

∂f/∂x = [∂f/∂x₁, ..., ∂f/∂xₙ]. Key identities: ∂(aᵀx)/∂x = a. ∂(xᵀx)/∂x = 2x. ∂(xᵀAx)/∂x = (A + Aᵀ)x = 2Ax if A is symmetric. ∂(aᵀXb)/∂X = abᵀ. Example: f = ||y - Xw||² = (y-Xw)ᵀ(y-Xw). Expand: f = yᵀy - 2wᵀXᵀy + wᵀXᵀXw. ∂f/∂w = -2Xᵀy + 2XᵀXw. Setting to zero: XᵀXw = Xᵀy — the normal equations. Derivation complete.

The trace trick

For a scalar f expressed as a trace: f = tr(AᵀB). Then ∂f/∂A = B. Since a scalar equals its own trace (tr(a) = a), we can write any scalar as a trace and use this identity. Useful when the scalar is expressed as a product of matrices: tr(AᵀBCD) can be differentiated by cycling the trace and applying the identity.

Vector by vector: the Jacobian

For y = f(x) with y ∈ ℝᵐ and x ∈ ℝⁿ, the Jacobian J ∈ ℝᵐˣⁿ has J_{ij} = ∂yᵢ/∂xⱼ. Key Jacobians: ∂(Ax)/∂x = A (linear map y=Ax: Jacobian is just A). ∂(xᵀA)/∂x = Aᵀ. ∂σ(x)/∂x = diag(σ(x) ⊙ (1-σ(x))) for element-wise sigmoid. ∂ReLU(x)/∂x = diag(𝟙[x > 0]) (diagonal of indicators). ∂softmax(x)/∂x = diag(s) - ssᵀ where s = softmax(x). The softmax Jacobian is not diagonal — each output depends on all inputs.

Backprop through a linear layer: full derivation

Forward pass: z = Wx + b, a = σ(z), where W ∈ ℝᵐˣⁿ, x ∈ ℝⁿ, b ∈ ℝᵐ. Suppose upstream gradient ∂L/∂a ∈ ℝᵐ is known. We need: ∂L/∂W, ∂L/∂b, ∂L/∂x (to pass backward). Step 1: ∂L/∂z = ∂L/∂a ⊙ σ'(z) (element-wise multiply by activation derivative). Let δ = ∂L/∂z ∈ ℝᵐ. Step 2 (gradient w.r.t. weights): ∂L/∂W_{ij} = Σₖ (∂L/∂zₖ)(∂zₖ/∂W_{ij}) = δᵢ xⱼ. Therefore: ∂L/∂W = δ xᵀ ∈ ℝᵐˣⁿ. The weight gradient is the outer product of the upstream gradient and the input. Step 3 (gradient w.r.t. bias): ∂L/∂b = δ. Step 4 (gradient w.r.t. input — to pass to previous layer): ∂L/∂x = Wᵀ δ ∈ ℝⁿ. Summary: δ_prev = Wᵀ δ (backprop through W is multiplication by Wᵀ). ∂L/∂W = δ xᵀ (outer product). This is the complete backprop update for a linear layer. Every framework implements exactly this.

Deriving the attention gradient (sketch)

Attention: Attention(Q, K, V) = softmax(QKᵀ/√d) V. Let S = QKᵀ/√d, A = softmax(S), output = AV. For scalar loss L and upstream gradient ∂L/∂(AV): ∂L/∂A = (∂L/∂(AV)) Vᵀ. ∂L/∂V = Aᵀ (∂L/∂(AV)). ∂L/∂S uses the softmax Jacobian (full matrix). ∂L/∂Q = (∂L/∂S) K / √d. ∂L/∂K = (∂L/∂S)ᵀ Q / √d. This is what PyTorch's autograd computes — the same chain of matrix products.

The Frobenius inner product

For matrix-valued functions, the gradient can be derived using the Frobenius inner product: ⟨A, B⟩_F = tr(AᵀB). The directional derivative of f(X) in direction dX is df = ⟨∇_X f, dX⟩_F = tr((∇_X f)ᵀ dX). To find ∇_X f: compute df, express it as tr(AᵀdX), then ∇_X f = A. Example: f = tr(AXB), df = tr(A dX B) = tr(BAᵀ ... ) ... = tr((AᵀB)ᵀ dX). So ∇_X f = (AᵀB)ᵀ = BᵀA.

Interview questions on this topic

"Derive ∂L/∂W for a linear layer given upstream gradient δ." — z = Wx + b, upstream δ = ∂L/∂z. ∂L/∂W_{ij} = Σₖ δₖ ∂zₖ/∂W_{ij} = δᵢ xⱼ. In matrix form: ∂L/∂W = δxᵀ. This is the outer product of the upstream gradient and the layer input.

"Why is the backward pass through a linear layer multiplication by Wᵀ?" — Forward: z = Wx (y = f(x) = Wx, Jacobian = W). Chain rule: ∂L/∂x = Wᵀ ∂L/∂z. The transpose appears because in numerator layout, the Jacobian of Wx w.r.t. x is W, and the chain rule for ∂L/∂x requires multiplying by the transpose Jacobian: ∂L/∂x = (∂z/∂x)ᵀ ∂L/∂z = Wᵀ δ.

"What is the gradient of the cross-entropy loss with softmax output?" — Let z be the logits, s = softmax(z), L = -Σ_k y_k log s_k. Using the softmax Jacobian: ∂L/∂z_i = s_i - y_i. The gradient is simply prediction minus label. This is one of the cleanest gradients in all of deep learning and is the reason softmax + cross-entropy is so well-behaved numerically.

"What shape is ∂L/∂W for a linear layer W ∈ ℝᵐˣⁿ and why?" — ∂L/∂W ∈ ℝᵐˣⁿ — the same shape as W. The gradient of a scalar with respect to a matrix has the same shape as the matrix. Element-wise: (∂L/∂W)_{ij} = ∂L/∂W_{ij} — the partial derivative of L with respect to each parameter.

Try on Colab: implement a single linear layer with manual forward and backward passes (no autograd). For an input x ∈ ℝ^{10}, W ∈ ℝ^{5×10}, and MSE loss on a target y ∈ ℝ^5, compute: (1) the forward pass z = Wx, (2) the loss L = ||z - y||², (3) δ = ∂L/∂z = 2(z-y), (4) ∂L/∂W = δxᵀ, (5) ∂L/∂x = Wᵀδ. Verify each gradient against PyTorch autograd using torch.autograd.gradcheck or finite differences.