HVP via grad-of-grad

How a gradient changes

A gradient measures how a scalar loss changes with its inputs. The Hessian measures how that gradient changes. For d inputs it is a d by d matrix of second derivatives.

Often we only need its product with a vector, H @ v: the rate at which the gradient changes along x + epsilon*v. This avoids constructing a dense Hessian with d*d entries. It does not guarantee total memory use of O(d), since the differentiated computation can store other arrays.

Make a scalar before differentiating again

jax.grad needs a scalar output. If x is a vector, jax.grad(f)(x) is also a vector, so calling jax.grad(jax.grad(f))(x) raises an error. It does not return a Hessian. jax.hessian(f)(x) would construct the full matrix, but that is unnecessary for a single product.

Instead, take the dot product of the gradient with a fixed v, then differentiate that scalar:

import jax
import jax.numpy as jnp

A = jnp.array([[2.0, 1.0], [1.0, 3.0]])
x = jnp.array([1.0, 2.0])
v = jnp.array([2.0, -1.0])

def f(x):
    return 0.5 * x @ A @ x

grad_f = jax.grad(f)
hvp = jax.grad(lambda z: jnp.vdot(grad_f(z), v))(x)
# hvp = [3.0, -1.0]

Differentiating this dot product gives H.T @ v. For a twice continuously differentiable real scalar function, the Hessian is symmetric, so this equals H @ v. Here A is symmetric and the loss is quadratic: its gradient is A @ x and its Hessian is A. Thus A @ v = [2*2 - 1, 1*2 - 3] = [3, -1].

For these real vectors, jnp.vdot is the usual dot product. Keep v fixed during differentiation. Capturing it in the outer scope, as above, is one way to do that. This construction uses reverse-mode differentiation through another reverse-mode derivative.

Compute derivatives by following the operations

Automatic differentiation computes derivatives by applying the chain rule to a program's operations. It does not estimate them by evaluating nearby inputs, as finite differences do, or rearrange the whole expression using symbolic algebra.

For f(g(x)), the chain rule combines the derivative of g with the derivative of f at g(x). Repeating this process lets a framework differentiate a long sequence of operations.

Forward mode follows an input change

For a function with n inputs and m outputs, the Jacobian J is an (m, n) matrix of derivatives. Entry (i, j) describes how output i changes with input j.

Forward-mode differentiation computes Jv: the output's rate of change when the input moves in direction v. For example, if f(x) = sum(x_i^2), then at x = [1, 2] the gradient is [2, 4]. In direction v = [3, 0], the rate of change is 2 * 3 + 4 * 0 = 6.

This is a directional derivative. It gives one Jacobian column only when v selects one input coordinate, such as [1, 0].

Reverse mode works backward from an output

Reverse mode computes J^T u, where u weights the outputs. For a scalar loss, choosing u = 1 gives the gradient with respect to all inputs. This makes reverse mode useful for training: a model can have many parameters but just one loss.

Forward mode can be useful when you need sensitivity to only a few input directions. Neither mode needs to construct the full Jacobian for these products. That can save substantial memory. Higher-order differentiation can similarly compute a Hessian-vector product without storing the full matrix of second derivatives.

Changing gradient flow changes what is differentiated

A stop-gradient operation keeps a value in the forward calculation but treats it as constant during differentiation. A custom derivative supplies a derivative rule explicitly. It should match the intended function unless the exercise asks for a surrogate, such as an approximation used during training.