Unit 2.5 · about 4 h
63 / 324Forward mode, reverse mode, JVPs, VJPs and Hessian-vector products
One reverse pass costs about one forward pass while forward mode costs one pass per input. That cost model decides how every large model is differentiated.
- derive
- implement
You can now
- Explain forward mode (JVP) and reverse mode (VJP) and choose between them from input and output counts.
- Implement forward-mode dual numbers and a reverse-mode accumulation for a small graph without forming any Jacobian.
- Compute a Hessian-vector product by forward-over-reverse or double backward and check it against a finite difference of gradients.
Covers
- Computational graphs and accumulation through branches
- Forward mode, reverse mode, JVPs and VJPs
- Hessian-vector products
Study, in this order
4 steps
- Forward and reverse-mode automatic differentiationMIT 18.S096 - Matrix Calculus for ML
Lecture 5, Part 3 notes (course notes Chapter 8, Sections 8.1 and 8.4): forward mode via dual numbers with the Babylonian square-root example, the cost rule (forward scales with inputs, reverse with outputs) and forward-over-reverse Hessian-vector products (Section 8.4.1)
- Dual numbers and both AD enginesHands-On ML - Autodiff notebook
Cells Forward mode autodiff using dual numbers (
DualNumberclass; seedingxgives 24.0 and seedingygives 10.0 forfat (3, 4)), Reverse mode autodiff (toy engine withbackpropagateon the samef), Reverse mode autodiff using PyTorch (torch.autograd.gradwithcreate_graph=Truefor second order)Depth · Rebuild the dual-number and reverse engines independently in your own file; do not copy the notebook cells.
- Backpropagation cost and graph equationsMathematics for Machine Learning - Ch. 5
Section 5.6 (pp. 159-164): explicit gradient (5.109)-(5.110) motivating the cost worry, forward and backward graph computation with accumulation over parents (5.135)-(5.145), reverse mode cheaper when inputs outnumber labels
- Reverse-mode graph, accumulation and grad modesPyTorch Autograd mechanics
How autograd encodes the history (reverse-mode graph recording); accumulation into
.gradandtorch.autograd.gradas a vector-Jacobian product; Grad Modes with Default, No-grad and Inference behaviorDepth · Consult after your handwritten engine. Skip multithreading and in-place correctness checks.
Focus on
- A JVP pushes a tangent forward (\(Jv\), one column); a VJP pulls a cotangent back (\(v^\top J\), one row). Forward cost scales with \(n\) inputs, reverse with \(m\) outputs.
- Reverse mode must keep or recompute forward values. That memory-for-time trade is the price of the single pass.
- Dual numbers evaluate \(f(x + \epsilon) = f(x) + f'(x) \epsilon\) exactly since \(\epsilon^2 = 0\); seeding each input in turn gives one gradient column per pass.
- An HVP is the JVP of the gradient. It needs \(C^2\) so the differentiation order does not matter.
Reproduce closed-book
DualNumber type (\(\epsilon^2 = 0\)) and recover \(df/dx = 24.0\) and \(df/dy = 10.0\) for \(f(x, y) = x^2 y + y + 2\) at (3, 4), matching the notebook analytic values. State the cost rule: \(n\) inputs with 1 output take 1 reverse pass versus \(n\) forward passes. Verify an HVP without forming \(H\) on \(f(x) = 0.5 x^\top Ax + b^\top x\) in \(\mathbb{R}^3\) with \(A = M + M^\top + 3I\), drawing \(M\) (\(3 \times 3\)), \(b\), \(x\) and \(v\) in that order with standard_normal from numpy.random.default_rng(11) in float64: the centered difference \((\operatorname{grad} f(x + hv) - \operatorname{grad} f(x - hv))/(2h)\) at \(h = 10^{-5}\) misses \(Av\) by 2.2e-11 in the 2-norm. The gradient is linear, so that gap is rounding alone.Traps
- Building the full Jacobian to get one JVP or VJP. Propagate the vector instead; the matrix is never needed.
- Calling
torch.autograd.functional.jacobianfor a check on wide maps. It materializes the matrix; usegrad(a VJP) or directional differences. - Assuming reverse mode always wins. For \(\mathbb{R}^1 \to \mathbb{R}^N\), one forward pass beats \(N\) reverse passes, so count outputs before choosing.
Practice
MIT 18.S096 Problem Set 2, Problem 1(c) (adjoint gradient plus finite-difference test)
Problem 1(c): implement the two-solve adjoint gradient procedure and test \(\operatorname{grad} f \cdot dp\) against \(f(p + dp) - f(p)\) on random \(a\), \(b\), \(c\), \(p\). Official solutions exist; attempt before reading them
Parts (a) and (b) derive the adjoint form; part (c) is the implementation. Record the generator, seed and dtype.
Self-test
-
Compare forward and reverse cost for \(f: \mathbb{R}^{1000} \to \mathbb{R}\).
Reverse costs about one forward pass; forward costs about 1000 passes, one per input. Uses: each primitive JVP or VJP costs \(O(1)\) forward equivalents. -
What do a JVP and a VJP compute, and which one does
torch.autograd.gradperform?JVP computes \(Jv\) (pushforward of a tangent); VJP computes \(v^\top J\) (pullback of a cotangent).torch.autograd.gradperforms a VJP. -
Give the two HVP routes and the assumption each needs.
Forward-over-reverse (JVP of the gradient) or double backward (VJP of the gradient). Both need \(C^2\) so mixed partials agree and the order is irrelevant. -
Why does reverse mode store forward values?
Local pullbacks need operand values (for example \(d(ab)\) needs \(a\) and \(b\)). The forward pass saves them; recomputation trades time for memory.
Completion criterion
Stop when
Your dual-number gradient matches (24.0, 10.0), your HVP matches a finite difference of gradients to 1e-9, and you pick the AD mode from \((n, m)\) without notes.
Leads to
- 2.6 Backpropagation through dense layers and gradient checks Dense-layer backprop is reverse mode at scale, and gradient checks test it.
- 1.4 Branches, broadcasting and saved values in backward Broadcast backward sums cotangents, which is the VJP of a broadcast; read it as this unit's pullback.
Comes after
- 2.2 Jacobians, Hessians and Taylor approximationsThe Hessian returns as Hessian-vector products computed without forming the matrix.
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.
- 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.
- Mathematics for Machine Learning - Ch. 5 Use Chapter 5 as the compact derivative, Jacobian, chain-rule and backpropagation reference.
Also in this section
- 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.