Unit 2.3 · about 4 h
61 / 324The chain rule and matrix derivatives
Matrix backprop is the chain rule with shapes. Differentials keep the product order that index notation hides, and that order decides the flop count.
- derive
You can now
- Differentiate composite scalar, vector and matrix functions with differentials and the operator chain rule.
- Derive \(d(X^{-1}) = -X^{-1} dX X^{-1}\) and gradients of affine, quadratic and trace objectives with correct shapes.
- Predict Jacobian accumulation cost for a chain of maps and choose left-to-right versus right-to-left order.
Covers
- Scalar and multivariable chain rule
- Matrix derivatives for affine, quadratic and least-squares objectives
Study, in this order
5 steps
- Operator chain rule and the inverse differentialMIT 18.S096 - Matrix Calculus for ML
Lecture 1, Part 2 notes (course notes Chapter 2, Sections 2.4-2.6): operator sum, product and chain rules, \(d(A^3) = A^2 dA + A dA A + dA A^2\), \(d(A^{-1}) = -A^{-1} dA A^{-1}\), forward versus reverse Jacobian accumulation
- Jacobians of matrix functionsMIT 18.S096 - Matrix Calculus for ML
Lecture 2, Part 2 notes (course notes Chapter 3): the Jacobian of a matrix function via \(\operatorname{vec}\), the \(4 \times 4\) Jacobian of matrix squaring, the identity \((A \otimes B) \operatorname{vec}(C) = \operatorname{vec}(BCA^\top)\), and why forming Kronecker Jacobians costs \(m^4\) (Section 3.3.3)
Depth · Read for the vec picture only; compute gradients with the Frobenius route in the next item.
- Gradients in general vector spacesMIT 18.S096 - Matrix Calculus for ML
Lecture 4, Part 1 notes (course notes Section 5.1): the Frobenius inner product \(\operatorname{tr}(A^\top B)\), gradients of matrix-input functions read off \(df = \operatorname{tr}(G^\top dA)\), with the Frobenius-norm and \(x^\top Ay\) (gradient \(xy^\top\)) examples
- Matrix gradients and gradient identitiesMathematics for Machine Learning - Ch. 5
Sections 5.2.2 and 5.4-5.5 (pp. 148-149, 155-158): chain rule as matrix multiplication, gradients of matrices, identities (5.99)-(5.108) covering transpose, trace, determinant, inverse, \(a^\top Xb\) and \(x^\top Bx\)
- Derivative of the determinant and inverseMIT 18.S096 - Matrix Calculus for ML
Lecture 5, Part 1 (course notes Chapter 7): derivative of the matrix determinant and inverse, including \(d \log \det X = \operatorname{tr}(X^{-1} dX)\)
Focus on
- \(d(fg) = (df)g + f(dg)\) keeps factor order. Matrices do not commute, so \(d(A^3)\) has three distinct terms.
- \(d(X^{-1})\) comes from differentiating the constant \(XX^{-1} = I\); the final left-multiplication by \(X^{-1}\) needs \(X\) to be invertible.
- Read gradients off differentials with the Frobenius product: \(df = \operatorname{tr}(G^\top dX)\) means \(\operatorname{grad} f = G\).
- For three or more composed maps, accumulate Jacobians from the small side. Scalar-output chains go left to right.
Reproduce closed-book
float64 NumPy with \(X = \left[\begin{smallmatrix}3&1\\1&2\end{smallmatrix}\right], V = \left[\begin{smallmatrix}0.5&-0.25\\0.25&0.5\end{smallmatrix}\right], h = 10^{-5}\): \(\|(X + hV)^{-1} - (X^{-1} - h X^{-1} V X^{-1})\|_F\) = 9.3e-12, a scaled residual of 0.093, consistent with an \(O(h^2)\) remainder.Traps
- Writing \(d(A^3) = 3 A^2 dA\). By hand with \(A = \left[\begin{smallmatrix}0&1\\0&0\end{smallmatrix}\right]\) and \(dA = \left[\begin{smallmatrix}0&0\\1&0\end{smallmatrix}\right]\): \(A^2 = 0\) so \(3 A^2 dA = 0\), but the product rule gives \(A dA A = A\), which is nonzero.
- Forming a vec-Kronecker Jacobian to get a gradient. The differential plus the Frobenius product gives the gradient with no \(n^2 \times n^2\) matrix.
- Quoting \(d \log \det X\) without invertibility. The formula needs invertible \(X\); the det gradient \(\det(X)(X^{-1})^\top\) shows the same requirement.
Practice
MIT 18.S096 Problem Set 1, Problems 1 and 2 (operator and matrix derivatives)
Problem 1 (linear-operation derivatives including transpose and trace) and Problem 2 (gradients, Jacobians and linear operators, including the \(\Theta(n^2)\) Sherman-Morrison argument). Official solutions exist; attempt before reading them
Write every answer as \(f'(x)[dx]\) first, then convert to the requested Jacobian or gradient form.
Self-test
-
Derive \(d(X^{-1})\) and name the step where invertibility is used.
\(XX^{-1} = I\) is constant, so \((dX)X^{-1} + X d(X^{-1}) = 0\) by the product rule. Left-multiplying by \(X^{-1}\) needs \(X\) invertible, giving \(d(X^{-1}) = -X^{-1} dX X^{-1}\). -
For \(f(X) = \operatorname{tr}(AXB)\), give \(\operatorname{grad} f\) and the cyclic step.
\(df = \operatorname{tr}(A dX B) = \operatorname{tr}(BA dX)\) by \(\operatorname{tr}(ABC) = \operatorname{tr}(CAB)\), so \(\operatorname{grad} f = (BA)^\top = A^\top B^\top\) under the Frobenius product. -
Chain \(\mathbb{R}^{100} \to \mathbb{R}^{50} \to \mathbb{R}^{20} \to \mathbb{R}^1\): which accumulation order, with flop counts?
Left to right (reverse mode). Jacobians are \(1 \times 20\), \(20 \times 50\), \(50 \times 100\): left-to-right costs \(1 \times 20 \times 50 + 1 \times 50 \times 100 = 6000\) multiplications, right-to-left costs \(20 \times 50 \times 100 + 1 \times 20 \times 100 = 102000\). Reverse wins by about 17x. -
Problem Set 1, Problem 2(b)(ii) computes a Jacobian in \(\Theta(n^2)\) given \(A^{-1}\). What is the structural reason?
With \(M = A + yx^\top\), \(df = -M^{-1} y (dx^\top f)\), so the Jacobian is the rank-1 outer product \(-(M^{-1} y) f(x)^\top\). Sherman-Morrison gives \(M^{-1} v = A^{-1} v - A^{-1} y (x^\top A^{-1} v)/(1 + x^\top A^{-1} y)\), \(\Theta(n^2)\) per vector once \(A^{-1}\) is known, so \(f(x)\), \(M^{-1} y\) and their outer product each cost \(\Theta(n^2)\).
Completion criterion
Stop when
You derive \(d(X^{-1})\) and one trace or quadratic gradient closed-book with shapes stated, and a NumPy finite-difference check agrees to 1e-9 relative.
Mark passed, then open 2.4 Only mark this when you can do it without notes.Leads to
- 2.4 Gradients of logistic, softmax cross-entropy and regularized losses Loss gradients apply these differentials to least squares, logistic and softmax objectives.
- 1.3 Computation graphs and reverse accumulation The scalar engine mechanizes this chain rule with topological backward; check it against centered differences.
Resources for this unit
- MIT 18.S096 - Matrix Calculus for ML Use lectures 1-5, the Hessian portion of lecture 7 and computational-graph portion of lecture 8; do the supplied problems and notebooks.
- Mathematics for Machine Learning - Ch. 5 Use Chapter 5 as the compact derivative, Jacobian, chain-rule and backpropagation reference.
Also in this section
- Hands-On ML - Autodiff notebook Work through numerical differentiation and both AD modes, then rebuild the essential reverse-mode engine independently.
- PyTorch Autograd mechanics Consult after the handwritten engine for accumulation, graph behavior, disabled tracking and vector-Jacobian products.
- MIT 18.02SC - Multivariable Calculus Unit 2 Use sessions 24-38 only if partial derivatives, gradients, Hessians or the multivariable chain rule are weak.
- CS231n - Linear Classification Read the Softmax classifier section for the cross-entropy loss formula with both the information-theory and MLE views, and Practical issues: Numeric stability for the max-subtraction trick. Derive the softmax gradient by hand; do not implement the course assignment here.