Hessian-vector products, Gauss-Newton and trace estimation

Before this lesson: lesson 03 (Jacobian, Hessian, symmetry of second derivatives), lesson 05 (JVP, VJP, forward and reverse mode, dual numbers), lesson 06 (softmax cross-entropy) and lesson 07 (step size of finite differences). From linear algebra: lesson 06 (eigenvalues of symmetric matrices, Rayleigh quotient), lesson 07 (positive semidefinite matrices) and lesson 14 (power iteration, Lanczos, conjugate gradient). Time: about 70 minutes, plus the notebook.

---

Why this matters #

A network with $n = 10^7$ parameters has a Hessian with $n^2 = 10^{14}$ entries: 800 TB in float64. You can never form it. Yet many tools need curvature:

• Sharpness. The largest Hessian eigenvalue $\lambda_{\max}$ limits the stable learning rate of gradient descent. Power iteration finds it with Hessian-vector products only.
• Second-order training. Newton-CG and Hessian-free optimization (Martens 2010) solve $Hd = -g$ or $Gd = -g$ with conjugate gradient. CG needs only products with the matrix.
• Influence functions and implicit gradients need $H^{-1}g$, again solved with products (lesson 09).
• Laplace approximations (lesson 11) use the Gauss-Newton matrix as the posterior precision.
• Trace estimation. Continuous normalizing flows need $\operatorname{tr}(\partial f/\partial z)$; a physics-informed loss with a Laplacian needs $\operatorname{tr}\nabla^2 u$. Hutchinson's estimator gives both from a few products.

All of these work with the matrix only through products. This lesson shows how to compute the products.

1. Predict first #

2. The picture #

Curvature is the rate of change of the gradient #

In lesson 05 you met the two products that automatic differentiation gives for a map $F$ with Jacobian $J$: the JVP $v \mapsto Jv$ (forward mode) and the VJP $u \mapsto J^\top u$ (reverse mode). Each costs a small constant times one evaluation of $F$, and neither forms $J$. This lesson applies them to one special map: the gradient.

Start in one variable. The second derivative tells how fast the slope changes:

In $n$ variables, the gradient $g(x) = \nabla f(x)$ is itself a map from $\R^n$ to $\R^n$. Its Jacobian is the Hessian $H = \nabla^2 f(x)$ (lesson 03). So walk along the line $x + tv$ and watch the gradient arrow. The velocity of its tip at $t = 0$ is

This is the Hessian-vector product (HVP): one column of information about $H$, in the direction you choose. A central difference of two gradients approximates it: $\big(\nabla f(x + \eta v) - \nabla f(x - \eta v)\big)/(2\eta)$.

Interactive figure: loading…

Two ways to differentiate a gradient #

Automatic differentiation can differentiate any program, and the gradient computation is a program. So you can run a second pass of automatic differentiation over the first one:

• Forward-over-reverse. Run the gradient program (a forward pass and a backward pass) with dual numbers, that is, with a tangent attached to every value. The tangent of the output $\nabla f$ is $Hv$. This is a JVP of the gradient.
• Reverse-over-reverse. Form the scalar $h(x) = \nabla f(x)^\top v$ and take its gradient with a second reverse pass. Then $\nabla h = Hv$. This is a VJP of the gradient.

The figure in section 3.1 runs both on a small example, row by row.

Gauss-Newton: keep the curvature of the loss, drop the curvature of the model #

Many losses have the form $L(\theta) = \ell(r(\theta))$: a model or residual map $r$ followed by a simple loss $\ell$, for example $L = \frac12\|r(\theta)\|^2$ in least squares. Replace $r$ by its linear approximation $r(\theta) + J\delta$, where $J$ is the Jacobian of $r$. The model $\frac12\|r + J\delta\|^2$ is a convex quadratic in the step $\delta$, with Hessian $J^\top J$. This is the Gauss-Newton matrix. The true Hessian of $L$ has one more term, which comes from the curvature of $r$ and is weighted by the residuals. That extra term can be negative. The Gauss-Newton matrix cannot.

Interactive figure: loading…

A trace from random signs #

Take a random vector $z$ with independent entries $+1$ or $-1$ (each with probability ½). Then

because $z_i^2 = 1$. Each cross term $z_i z_j$ is $+1$ or $-1$ with equal chance, so it has mean 0. The average of $z^\top A z$ over many probes is therefore $\operatorname{tr} A$. Each probe needs only one product $Az$; for $A = H$ that is one HVP.

3. Definitions and main results #

Notation. $f: \R^n \to \R$ is twice continuously differentiable ($C^2$). Its gradient is the column $g(x) = \nabla f(x) \in \R^n$ and its Hessian is $H = \nabla^2 f(x) \in \R^{n\times n}$, with $H_{ij} = \partial^2 f/\partial x_i \partial x_j$. The vector $v \in \R^n$ is fixed: it does not depend on $x$.

3.1 Hessian-vector products #

Proposition 1 (two exact HVPs). (a) $\displaystyle \frac{d}{dt}\, g(x + tv)\Big|_{t=0} = Hv$ (forward-over-reverse: a JVP of $g$). (b) With $h(x) = g(x)^\top v$, $\nabla h(x) = Hv$ (reverse-over-reverse: a VJP of $g$).

Proof.

1. The Jacobian of $g$ is the Hessian: $[Dg(x)]_{ij} = \partial_j g_i = \partial_j \partial_i f = H_{ij}$.
2. (a) The derivative of $t \mapsto g(x + tv)$ at $t = 0$ is $Dg(x)\, v$ (chain rule, lesson 04). By step 1 this is $Hv$.
3. (b) $h(x) = v^\top g(x)$ with $v$ constant, so $\nabla h(x) = Dg(x)^\top v = H^\top v$.
4. $f$ is $C^2$, so $H$ is symmetric (Schwarz's theorem, lesson 03) and $H^\top v = Hv$. $\blacksquare$

Step 4 is where the Hessian is special. For a map $g$ that is not a gradient, the JVP gives $Dg\,v$ and the VJP gives $Dg^\top v$, and the two differ.

There are four ways to nest the two modes (Blondel and Roulet, §9.1):

Method	What it computes	Time	Memory
forward-over-reverse	JVP of the gradient program	small constant × one gradient	like one gradient, plus the tangents
reverse-over-reverse	gradient of $h = g^\top v$	small constant × one gradient	more: it also stores the first backward pass
reverse-over-forward	gradient of $x \mapsto Df(x)\,v$ (the same $h$)	small constant × one gradient	like forward-over-reverse
forward-over-forward	second-order JVPs, one per basis direction	$n$ passes for one $Hv$	small

Forward-over-reverse is the usual default (for example jax.jvp(jax.grad(f), ...)). Reverse-over-reverse is the "double backward" of PyTorch: it is available whenever the backward pass is itself differentiable.

Interactive figure: loading…

Why one HVP costs a small constant times one gradient.

1. Reverse mode gives $\nabla f$ for a small constant times the cost of $f$, independent of $n$. This is the cheap gradient principle (lesson 05).
2. The gradient computation is itself a program $g$ made of elementary operations. Its cost is $c \cdot \operatorname{cost}(f)$ for a small constant $c$.
3. Forward mode on any program costs a constant factor times that program: each elementary operation gets one extra derivative operation.
4. So the JVP of $g$ costs a constant times $\operatorname{cost}(g)$: a small multiple of one gradient. In practice this is about 2 to 4 gradients. No factor $n$ appears.
5. Reverse-over-reverse: $h = g^\top v$ costs $\operatorname{cost}(g) + n$, and reverse mode on $h$ costs a constant times that. The time is of the same order. But the second reverse sweep must also store the intermediate values of the first backward pass, so it needs more memory than forward-over-reverse.

The finite-difference HVP. The central difference $\big(g(x + \eta v) - g(x - \eta v)\big)/(2\eta)$ costs two gradients, but it is not exact. Its truncation error is $O(\eta^2)$ (from third derivatives of $f$) and its rounding error is $O(\varepsilon/\eta)$, where $\varepsilon \approx 2.2\times10^{-16}$ is float64 machine epsilon. As in lesson 07, the best step is about $\varepsilon^{1/3} \approx 6\times10^{-6}$ (for a well-scaled problem) and the best relative error is about $\varepsilon^{2/3} \approx 4\times10^{-11}$. In float32 ($\varepsilon \approx 1.2\times10^{-7}$) the best is only about $2\times10^{-5}$.

The full Hessian costs $n$ HVPs, with $v = e_1, \dots, e_n$: the HVP with $e_j$ is column $j$ of $H$. This is fine for $n = 10^3$ and impossible for $n = 10^7$.

3.2 The Gauss-Newton matrix #

Let $L(\theta) = \ell(r(\theta))$ with a model or residual map $r: \R^n \to \R^m$, its Jacobian $J = Dr(\theta) \in \R^{m\times n}$, and a $C^2$ loss $\ell: \R^m \to \R$. Write $r_i$ for the $i$-th output of $r$ and $\partial_i \ell$ for $\partial \ell/\partial r_i$.

Proposition 2 (Hessian of a composition).

Proof.

1. Chain rule: $\nabla L = J^\top \nabla \ell(r)$. In coordinates, $\partial_a L = \sum_i \partial_i \ell\; \partial_a r_i$.
2. Differentiate with respect to $\theta_b$ with the product rule. The first factor $\partial_i\ell$ depends on $\theta$ through $r$, so $\partial_b(\partial_i \ell) = \sum_j \partial_{ij}\ell\; \partial_b r_j$:

3. The first sum is entry $(a, b)$ of $J^\top \nabla^2 \ell\, J$. The second is entry $(a, b)$ of $\sum_i \partial_i\ell\, \nabla^2 r_i$. $\blacksquare$

Definition. The (generalized) Gauss-Newton matrix is $G = J^\top \nabla^2\ell(r)\, J \in \R^{n\times n}$. For $\ell = \frac12\|r\|^2$ we have $\nabla^2 \ell = I$ and $G = J^\top J$.

Three properties.

1. $G$ is positive semidefinite when $\ell$ is convex. For any $v$, $v^\top G v = (Jv)^\top\, \nabla^2 \ell\, (Jv) \ge 0$, because $\nabla^2\ell \succeq 0$. This holds even where $\nabla^2 L$ is indefinite.
2. $G = \nabla^2 L$ when the residual term vanishes: when $r$ is linear in $\theta$ ($\nabla^2 r_i = 0$), or when $\nabla \ell(r) = 0$ (for least squares: a perfect fit, $r = 0$). Near a good fit, $G \approx \nabla^2 L$.
3. Matrix-free product. To compute $Gv$: (i) $u = Jv$, one JVP of $r$; (ii) $w = \nabla^2 \ell\; u$, which is cheap: for $\ell = \frac12\|r\|^2$, $w = u$; for softmax cross-entropy with probabilities $p$, $w = p \odot u - p\,(p^\top u)$; (iii) $Gv = J^\top w$, one VJP of $r$.

For softmax cross-entropy, $\nabla^2_z \ell = \operatorname{diag}(p) - pp^\top$ with respect to the logits $z$ (lesson 06). It is positive semidefinite: $u^\top(\operatorname{diag}(p) - pp^\top)u = \sum_i p_i u_i^2 - (\sum_i p_i u_i)^2$ is the variance of the numbers $u_i$ under the probabilities $p_i$.

The Gauss-Newton step solves $G\delta = -\nabla L$. If $G \succ 0$ and $\nabla L \ne 0$, then $\nabla L^\top \delta = -\delta^\top G \delta < 0$, so the step is a descent direction. The Levenberg-Marquardt method adds damping, $(G + \mu I)\delta = -\nabla L$ with $\mu > 0$, which keeps the system positive definite and shortens the step when the linear model is poor. The optimization track covers these methods in detail (optimization).

Gauss-Newton, Fisher information and natural gradient

Let $\ell(r) = -\log p(y \mid r)$ be a negative log-likelihood, where the model output $r$ is the natural parameter of an exponential family: the logits of a softmax, or the mean of a Gaussian with fixed variance. Then $\nabla_r^2 \ell$ does not depend on the label $y$, and $G$ equals the Fisher information matrix of the model (Martens 2020; Blondel and Roulet §9.3). This connects Gauss-Newton to the natural gradient and to K-FAC (Martens and Grosse 2015), which approximates $G$ by Kronecker factors, one pair per layer.

3.3 Hutchinson's trace estimator #

Proposition 3. Let $A \in \R^{n\times n}$, and let $z \in \R^n$ be random with $\mathbb E[zz^\top] = I$. Then $\mathbb E[z^\top A z] = \operatorname{tr} A$. If $A$ is symmetric and $z$ is Rademacher (independent entries, $+1$ or $-1$ with probability ½ each), then $\operatorname{Var}(z^\top A z) = 2\sum_{i \ne j} A_{ij}^2$.

Proof.

1. A number equals its own trace, so $z^\top A z = \operatorname{tr}(z^\top A z) = \operatorname{tr}(A\, zz^\top)$ (cyclic property of the trace).
2. Trace and expectation are both linear, so $\mathbb E[z^\top A z] = \operatorname{tr}(A\, \mathbb E[zz^\top]) = \operatorname{tr}(A\, I) = \operatorname{tr} A$.
3. For Rademacher $z$, $z_i^2 = 1$. With $A$ symmetric, $z^\top A z - \operatorname{tr} A = \sum_{i \ne j} A_{ij} z_i z_j = 2\sum_{i<j} A_{ij} z_i z_j$.
4. Each product $z_i z_j$ ($i < j$) has mean 0 and variance 1. Two different products are uncorrelated: for example $\mathbb E[z_1 z_2 \cdot z_1 z_3] = \mathbb E[z_2]\,\mathbb E[z_3] = 0$.
5. The variance of a sum of uncorrelated terms is the sum of the variances: $4\sum_{i<j} A_{ij}^2 = 2\sum_{i \ne j} A_{ij}^2$. $\blacksquare$

The Hutchinson estimator averages $S$ independent probes: $\widehat{\operatorname{tr}} = \frac1S \sum_{s=1}^S z_s^\top A z_s$. It needs $S$ products $Az_s$ (for $A = H$: $S$ HVPs), and its variance is the single-probe variance divided by $S$. So the standard error falls like $1/\sqrt S$: four times more probes halve the error.

Variance for general probes, and why Rademacher is best

Let the entries $z_i$ be independent with mean 0, variance 1 and fourth moment $\kappa = \mathbb E[z_i^4]$, and let $A$ be symmetric.

1. Split $z^\top A z = \sum_i A_{ii} z_i^2 + \sum_{i\ne j} A_{ij} z_i z_j$.
2. The diagonal part has variance $\sum_i A_{ii}^2 \operatorname{Var}(z_i^2) = (\kappa - 1)\sum_i A_{ii}^2$.
3. The off-diagonal part has variance $2\sum_{i \ne j} A_{ij}^2$, as in the proof above.
4. The two parts are uncorrelated: $\mathbb E[z_i^2 z_j z_k] = 0$ for $j \ne k$, because $z_j$ or $z_k$ appears with power 1 and has mean 0.
5. So $\operatorname{Var}(z^\top A z) = 2\sum_{i\ne j} A_{ij}^2 + (\kappa - 1)\sum_i A_{ii}^2$.

Gaussian probes have $\kappa = 3$, which gives $2\|A\|_F^2$. Rademacher probes have $\kappa = 1$, the smallest possible value (since $\mathbb E[z_i^4] \ge (\mathbb E[z_i^2])^2 = 1$), and they add nothing for the diagonal. For a non-symmetric $A$, use its symmetric part $\frac12(A + A^\top)$: it has the same quadratic form $z^\top A z$.

Interactive figure: loading…

3.4 Algorithms that need only $Hv$ #

• Largest eigenvalue (sharpness). Power iteration $v \leftarrow Hv / \|Hv\|$ converges to the eigenvector of the eigenvalue with the largest absolute value (linear-algebra lesson 05). The Rayleigh quotient $v^\top H v$ then estimates that eigenvalue. Near a minimum, $H \succeq 0$ and this is $\lambda_{\max}$. The Lanczos method (linear-algebra lesson 14) gets several extreme eigenvalues with fewer HVPs.
• Newton-CG. The Newton step solves $Hd = -g$. Conjugate gradient solves it with one HVP per iteration, and stopping after a few iterations gives a "truncated Newton" step. CG needs a positive definite matrix, which is one reason to use $G$ (or $G + \mu I$) instead of $H$.
• Trace and diagonal. Hutchinson's estimator gives $\operatorname{tr} H$. The same probes give an estimate of the diagonal: $\mathbb E[z \odot Hz] = \operatorname{diag}(H)$ for Rademacher $z$.

Answers to section 1. (1) About $3T$: a small constant times one gradient (in practice 2 to 4), never $nT$, by the cost argument in section 3.1. (2) No. For a squared-error loss, $v^\top G v = \|Jv\|^2 \ge 0$ for every $v$, so $G$ has no negative eigenvalue. The negative curvature lives in the residual term. (3) One. For a diagonal $A$, $z^\top A z = \sum_i A_{ii} z_i^2 = \operatorname{tr} A$ for every Rademacher probe, so the variance $2\sum_{i\ne j}A_{ij}^2$ is 0.

4. Worked example #

(a) An HVP both ways. $f(x) = x_1^3 + x_1 x_2^2$ at $x = (1, 2)$ with $v = (1, 1)$. The gradient is $g(x) = (3x_1^2 + x_2^2,\; 2x_1 x_2)$, so $g(1, 2) = (7, 4)$.

• Forward-over-reverse. Move along the line: $x + tv = (1 + t,\, 2 + t)$, so the two entries of $g(x + tv)$ are $3(1+t)^2 + (2+t)^2$ and $2(1+t)(2+t)$. Their derivatives in $t$ are $6(1+t) + 2(2+t)$ and $2(2+t) + 2(1+t)$. At $t = 0$ this is $(6 + 4,\; 4 + 2) = (10, 6)$.
• Reverse-over-reverse. $h(x) = g(x)^\top v = 3x_1^2 + x_2^2 + 2x_1x_2$. Its gradient is $\nabla h = (6x_1 + 2x_2,\; 2x_2 + 2x_1)$. At $(1, 2)$: $(6 + 4,\; 4 + 2) = (10, 6)$.
• Check with the full Hessian. $H = \begin{pmatrix} 6x_1 & 2x_2 \\ 2x_2 & 2x_1 \end{pmatrix} = \begin{pmatrix} 6 & 4 \\ 4 & 2 \end{pmatrix}$ and $Hv = (6 + 4,\; 4 + 2) = (10, 6)$. ✓
• $\det H = 12 - 16 = -4 < 0$, so the eigenvalues have opposite signs: the point is saddle-like.
• Finite difference. $g$ is quadratic in $t$, so the central difference $\big(g(x + \eta v) - g(x - \eta v)\big)/(2\eta)$ equals $(10, 6)$ for every $\eta$, up to rounding.

(b) Gauss-Newton against the Hessian. Residuals $r(\theta) = (\theta_1\theta_2 + 2,\; \theta_1 - 1)$, loss $L = \frac12\|r\|^2$, at $\theta = (1, 2)$, with $v = (1, -1)$.

1. $r = (1\cdot 2 + 2,\; 1 - 1) = (4, 0)$, and $J = \begin{pmatrix} \theta_2 & \theta_1 \\ 1 & 0 \end{pmatrix} = \begin{pmatrix} 2 & 1 \\ 1 & 0 \end{pmatrix}$.
2. Matrix-free product: $u = Jv = (2 - 1,\; 1) = (1, 1)$, then $Gv = J^\top u = (2 + 1,\; 1 + 0) = (3, 1)$.
3. The matrix itself: $G = J^\top J = \begin{pmatrix} 5 & 2 \\ 2 & 1 \end{pmatrix}$. Its eigenvalues are $3 \pm 2\sqrt2 \approx 5.83$ and $0.17$: positive definite.
4. Residual term: only $r_1$ is nonzero, and $\nabla^2 r_1 = \begin{pmatrix} 0 & 1 \\ 1 & 0 \end{pmatrix}$ (the second derivatives of $\theta_1\theta_2$). So the term is $4\begin{pmatrix} 0 & 1 \\ 1 & 0 \end{pmatrix}$.
5. $\nabla^2 L = \begin{pmatrix} 5 & 6 \\ 6 & 1 \end{pmatrix}$, with eigenvalues $3 \pm 2\sqrt{10} \approx 9.32$ and $-3.32$: indefinite. Also $\nabla^2 L\, v = (5 - 6,\; 6 - 1) = (-1, 5) \ne Gv$.
6. The steps. The gradient is $\nabla L = J^\top r = (8, 4)$. The Gauss-Newton step solves $G\delta = -(8, 4)$: $\delta = (0, -4)$. It lands at $\theta = (1, -2)$, where $r = (-2 + 2,\; 0) = (0, 0)$: an exact solution in one step. The Newton step $-(\nabla^2 L)^{-1}\nabla L \approx (-0.52, -0.90)$ also goes downhill here, but this is not guaranteed when $\nabla^2 L$ is indefinite.

The large residual $r_1 = 4$ makes the dropped term large, so $G$ and $\nabla^2 L$ differ a lot here.

(c) Hutchinson by hand. $A = \begin{pmatrix} 2 & 1 \\ 1 & 3 \end{pmatrix}$, $\operatorname{tr} A = 5$. For a Rademacher probe, $z^\top A z = 2z_1^2 + 3z_2^2 + 2z_1z_2 = 5 + 2z_1z_2$. The four possible probes $(\pm1, \pm1)$ give $7, 3, 3, 7$: mean 5 and variance $4 = 2(1^2 + 1^2)$. Gaussian probes would give variance $2\|A\|_F^2 = 2(4 + 1 + 1 + 9) = 30$.

5. In code #

NumPy checks the numbers of section 4. The gradients are written by hand.

import itertools
import numpy as np

g = lambda x: np.array([3 * x[0]**2 + x[1]**2, 2 * x[0] * x[1]])   # gradient of x1^3 + x1 x2^2
H = lambda x: np.array([[6 * x[0], 2 * x[1]], [2 * x[1], 2 * x[0]]])
x, v = np.array([1.0, 2.0]), np.array([1.0, 1.0])
for eta in (1e-1, 1e-5):
    print((g(x + eta * v) - g(x - eta * v)) / (2 * eta))   # [10. 6.] both times: g is quadratic
print(H(x) @ v)                                            # [10. 6.]

# Gauss-Newton product with one JVP (J @ v) and one VJP (J.T @ u)
r = lambda th: np.array([th[0] * th[1] + 2, th[0] - 1])
th, v = np.array([1.0, 2.0]), np.array([1.0, -1.0])
J = np.array([[th[1], th[0]], [1.0, 0.0]])
HL = J.T @ J + r(th)[0] * np.array([[0.0, 1.0], [1.0, 0.0]])
print(J.T @ (J @ v), HL @ v)                               # [3. 1.] [-1. 5.]
print(np.linalg.eigvalsh(J.T @ J), np.linalg.eigvalsh(HL)) # [0.17 5.83] [-3.32 9.32]
print(np.linalg.solve(J.T @ J, -J.T @ r(th)))              # [0. -4.]: lands on r = 0

# Hutchinson: all four Rademacher probes
A = np.array([[2.0, 1.0], [1.0, 3.0]])
q = [z @ A @ z for z in map(np.array, itertools.product([-1, 1], repeat=2))]
print(q, np.mean(q), np.var(q))                            # [7, 3, 3, 7] 5.0 4.0

In PyTorch, torch.func composes the two modes directly. The eager "double backward" needs create_graph=True in the first call, so that the gradient is itself recorded in a graph:

import torch
from torch.func import grad, jvp, vjp

torch.set_default_dtype(torch.float64)
f = lambda x: x[0]**3 + x[0] * x[1]**2
x, v = torch.tensor([1.0, 2.0]), torch.tensor([1.0, 1.0])

print(jvp(grad(f), (x,), (v,))[1])          # forward-over-reverse: tensor([10., 6.])
print(grad(lambda x: grad(f)(x) @ v)(x))    # reverse-over-reverse: tensor([10., 6.])

xr = x.clone().requires_grad_(True)         # eager double backward
(gx,) = torch.autograd.grad(f(xr), xr, create_graph=True)
print(torch.autograd.grad(gx @ v, xr)[0])   # tensor([10., 6.])

r = lambda th: torch.stack([th[0] * th[1] + 2, th[0] - 1])
th, v = torch.tensor([1.0, 2.0]), torch.tensor([1.0, -1.0])
u = jvp(r, (th,), (v,))[1]                  # u = J v      -> tensor([1., 1.])
print(vjp(r, th)[1](u)[0])                  # G v = J^T u  -> tensor([3., 1.])

In JAX the same forward-over-reverse HVP is jax.jvp(jax.grad(f), (x,), (v,))[1].

6. Where it shows up #

• Sharpness and the edge of stability. On a quadratic, gradient descent with step $\eta$ is stable only if $\eta < 2/\lambda_{\max}$. Cohen et al. (2021) observed that in full-batch training the sharpness $\lambda_{\max}(\nabla^2 L)$ rises until it reaches about $2/\eta$ and then stays there. Such studies measure $\lambda_{\max}$ with power iteration or Lanczos on HVPs; PyHessian (Yao et al. 2020) also estimates the trace with Hutchinson and the eigenvalue density with stochastic Lanczos.
• Second-order optimization. Hessian-free optimization (Martens 2010) runs CG on damped Gauss-Newton products $(G + \mu I)v$. K-FAC and Shampoo approximate $G$ or the Fisher matrix with Kronecker factors.
• Influence functions. The effect of one training point on a prediction needs $H^{-1}\nabla \ell$ (Koh and Liang 2017). It is computed with CG or a truncated Neumann series on HVPs.
• Laplace approximation (lesson 11). The common "GGN Laplace" posterior uses the covariance $(G + \text{prior precision})^{-1}$, because $G$ is positive semidefinite and $H$ may not be.
• Trace estimation in generative models. A continuous normalizing flow (FFJORD, Grathwohl et al. 2019) needs $\operatorname{tr}(\partial f/\partial z)$ in its log-density ODE. Hutchinson with one VJP per probe replaces $d$ backward passes by one. Sliced score matching uses the same trick (generative models).
• ML for science. In 4D-Var data assimilation for weather prediction, the "tangent-linear model" is the JVP and the "adjoint model" is the VJP of the forecast model. Incremental 4D-Var is a Gauss-Newton method, and its inner solver uses exactly $J^\top(Jv)$. Full-waveform inversion in seismology uses the same Gauss-Newton and truncated-Newton products. A physics-informed loss with a Laplacian needs $\Delta u = \operatorname{tr}\nabla_x^2 u$: $d$ HVPs for input dimension $d$, or a Hutchinson estimate when $d$ is large (ML for science).

7. Common mistakes #

• Forming the full matrix. If you need $Hv$, $Gv$ or $\operatorname{tr} H$, do not form $H$, $G$ or $J$. Use products.
• Letting $v$ depend on $x$ in reverse-over-reverse. If $v = v(x)$ is not detached from the graph, then $\nabla(g^\top v) = Hv + Dv(x)^\top g$: an extra term appears.
• Forgetting create_graph=True in the first backward call of a double-backward HVP. Without it the first gradient is not part of any graph, and the second call raises an error ("does not require grad").
• Treating $G$ as the Hessian. $G$ drops $\sum_i \partial_i\ell\, \nabla^2 r_i$. It is positive semidefinite by construction, so it cannot show negative curvature or saddle points. Far from a good fit it can be a poor model of $L$.
• Using power iteration as if it returned $\lambda_{\max}$. It returns the eigenvalue of largest absolute value. If $H$ has a large negative eigenvalue, that is what you get.
• Gaussian probes by default. They add $2\sum_i A_{ii}^2$ to the variance. Rademacher probes do not.
• Finite-difference HVPs in float32. With $\varepsilon \approx 1.2\times10^{-7}$ the best relative error is only about $2\times10^{-5}$. Use the exact autodiff HVP.

8. Check yourself #

1. Reverse-over-reverse computes $Dg(x)^\top v$. Why does this equal the HVP, and what goes wrong for a vector field $g$ that is not a gradient?
   Answer

   For $g = \nabla f$ with $f \in C^2$, $Dg = \nabla^2 f$ is symmetric (Schwarz), so $Dg^\top v = Dg\,v = Hv$. For a general $g$, $Dg$ is not symmetric, and the VJP $Dg^\top v$ differs from the JVP $Dg\,v$.
2. $r(\theta) = \theta^2 - 2$ (one residual, one parameter) and $L = \frac12 r^2$. Give $G$ and $\nabla^2 L$ at $\theta = 1$ and at $\theta = 0.5$.
   Answer

   $r' = 2\theta$ and $r'' = 2$, so $G = (r')^2 = 4\theta^2$ and $\nabla^2 L = 4\theta^2 + r \cdot 2 = 6\theta^2 - 4$. At $\theta = 1$: $G = 4$, $\nabla^2 L = 2$. At $\theta = 0.5$: $G = 1$, $\nabla^2 L = -2.5$. The true curvature is negative there; $G$ stays positive.
3. $A = \begin{pmatrix} 1 & 2 \\ 2 & 1 \end{pmatrix}$. What is the single-probe variance of Hutchinson's estimator with Gaussian probes? What is the standard deviation of the mean of 100 Rademacher probes?
   Answer

   Gaussian: $2\|A\|_F^2 = 2(1 + 4 + 4 + 1) = 20$. Rademacher: $2(2^2 + 2^2) = 16$ per probe, so the mean of 100 probes has standard deviation $\sqrt{16/100} = 0.4$, around the true trace 2.
4. For softmax cross-entropy with probabilities $p$, how do you apply $\nabla^2\ell$ to a vector $u$ inside a Gauss-Newton product, and why is $G$ positive semidefinite?
   Answer

   $\nabla^2_z \ell = \operatorname{diag}(p) - pp^\top$, so $w = p\odot u - p\,(p^\top u)$. And $u^\top(\operatorname{diag}(p) - pp^\top)u = \sum_i p_i u_i^2 - (\sum_i p_i u_i)^2$ is the variance of the numbers $u_i$ under $p$, so it is $\ge 0$. Hence $v^\top G v = (Jv)^\top \nabla^2\ell\,(Jv) \ge 0$.
5. Which HVP method needs less memory, forward-over-reverse or reverse-over-reverse, and why?
   Answer

   Forward-over-reverse. The tangents travel with the forward and backward passes and can be freed as the passes go. Reverse-over-reverse must also store the intermediate values of the first backward pass, because the second reverse sweep visits them.
6. You run power iteration with HVPs and get the estimate $-50$. What does this tell you about $H$?
   Answer

   The eigenvalue of largest absolute value is about $-50$: $H$ has strong negative curvature in some direction, so the point is not a minimum. All eigenvalues lie in about $[-50, 50]$, so the shifted matrix $H + 50I$ has eigenvalues $\ge 0$ (about). Power iteration on it finds $\lambda_{\max} + 50$; one HVP of $H + 50I$ is $Hv + 50v$. Lanczos gives both ends directly.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Write the Hessian-vector product as the derivative of the gradient.

   Show answer

2. How does reverse-over-reverse compute $Hv$?

   Show answer

3. Why do the JVP and the VJP of the gradient give the same vector?

   Show answer

4. What does one exact HVP cost, compared with one gradient?

   Show answer

5. How many HVPs do you need to form the full $n \times n$ Hessian?

   Show answer

6. What is the Gauss-Newton matrix of $L = \ell(r(\theta))$?

   Show answer

7. Why is the Gauss-Newton matrix positive semidefinite?

   Show answer

8. How do you compute $J^\top J v$ without forming $J$?

   Show answer

9. What is Hutchinson's estimator, and why is it unbiased?

   Show answer

10. Variance of one Rademacher Hutchinson probe for $A = \begin{pmatrix} 2 & 1 \\ 1 & 3\end{pmatrix}$?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "How do you get the top Hessian eigenvalue of a 1B-parameter model?" Power iteration or Lanczos on HVPs. Each HVP costs a few gradients; never form $H$. Power iteration returns the eigenvalue of largest absolute value.
• "What does an HVP cost, and how is it computed?" A small constant times a gradient. Forward-over-reverse: a JVP of the gradient function. Or double backward: the gradient of $\nabla f^\top v$.
• "Why is Gauss-Newton safer than Newton on a neural network?" $G = J^\top\nabla^2\ell\,J \succeq 0$ for convex $\ell$, so the quadratic model is convex and the step is a descent direction. $\nabla^2 L$ is indefinite in general.
• "How do you get the trace of a Jacobian in a continuous normalizing flow without forming it?" Hutchinson: $z^\top (\partial f/\partial z)\, z$ with one VJP per probe, Rademacher probes.
• "When is Gauss-Newton equal to the Hessian?" When the residual term vanishes: a linear model, or zero residuals at a perfect fit.

10. Practice #

The notebook for this lesson runs in Colab. It uses NumPy only: a small dual-number class (from lesson 05) gives forward-over-reverse on hand-written gradients. Five exercises, each with a check cell and a folded solution:

1. hvp_fwd_rev and hvp_fd on a logistic-regression loss; compare with the exact $X^\top S X v / B + \mu v$.
2. hessian_from_hvp: build $H$ column by column and check that it is symmetric.
3. gauss_newton_vp for a small exponential curve fit; compare $G$ with the full Hessian near and far from the fit.
4. hutchinson_trace with Rademacher and Gaussian probes; compare the empirical variance with the formula.
5. top_eigenvalue: power iteration with HVPs for $\lambda_{\max}$ of a small MLP loss Hessian.

11. Watch and read #

Lecture 7 Part 2: Second Derivatives, Bilinear Forms, and Hessian MatricesAlan Edelman and Steven G. Johnson · MIT 18.S096 Matrix Calculus (IAP 2023) · 46 min

The Hessian as the derivative of the gradient, and second derivatives as bilinear forms.

Lecture 4 - Automatic DifferentiationTianqi Chen and Zico Kolter · CMU Deep Learning Systems · 64 min

Reverse mode by extending the computational graph: this is why you can take the gradient of a gradient.

• Blondel and Roulet, The Elements of Differentiable Programming (free): §9.1 (the four HVP methods and their costs), §9.2 (Gauss-Newton matrix and Gauss-Newton vector product), §9.8.1 (Girard-Hutchinson trace estimator).
• JAX, The Autodiff Cookbook: sections on Hessian-vector products with grad-of-grad and with forward-over-reverse.
• PyTorch tutorial, Jacobians, Hessians, hvp, vhp, and more (torch.func).
• Pearlmutter (1994), Fast exact multiplication by the Hessian: the original exact HVP.
• Martens (2010), Deep learning via Hessian-free optimization: CG on Gauss-Newton products.
• Yao et al. (2020), PyHessian: Neural Networks Through the Lens of the Hessian: power iteration, Hutchinson and stochastic Lanczos on HVPs.
• Cohen et al. (2021), Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability.
• Bright, Edelman and Johnson, Matrix Calculus (for Machine Learning and Beyond): the course notes of the MIT lectures, with a chapter on second derivatives.