HVP via jvp-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.

Take a JVP of the gradient

jax.jvp returns a function's value and its directional derivative. Apply it to the gradient function: the value is the gradient, and the directional derivative is the Hessian-vector product.

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)
gradient, hvp = jax.jvp(grad_f, (x,), (v,))
# gradient = A @ x = [4.0, 7.0]
# hvp = A @ v = [3.0, -1.0]

The two results are different: the first describes the loss's slope at x, while the second describes how that slope changes along v. The Jacobian of grad_f is the Hessian, which explains why its JVP is H @ v. Creating grad_f constructs a callable; it does not yet evaluate a gradient at x.

This combines forward-mode differentiation with a reverse-mode gradient, often called forward-over-reverse. It can use less memory than nesting two reverse-mode derivatives, but it is not guaranteed to be faster for every function or backend. The two approaches compute the same mathematical product where the Hessian is symmetric; floating-point results can differ slightly.

Pass the single input and direction as (x,) and (v,). Matching lists are also supported. The real floating-point arrays x and v must have the same shape and dtype. Return the second result, not the gradient.

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.