Hessian Vector Products

Multiply by the Hessian of a function without ever building the Hessian.

For a scalar function f of a vector x with n entries, the Hessian H is the n x n matrix of second derivatives. For a network with a million parameters it has a trillion entries, so nobody stores it. But H @ v is cheap. It is the gradient of grad(f) . v:

g  = gradient of f at x          (keep the graph, so g can be differentiated)
Hv = gradient of (g . v) at x

Write two functions:

1. hvp(f, x, v) returns H @ v with two calls to torch.autograd.grad. Use create_graph=True on the first one.
2. top_hessian_eigenvalue(f, x, iters) estimates the eigenvalue of H with the largest absolute value by power iteration that uses only hvp. Start from v = ones(n) / sqrt(n), repeat iters times v = hvp(v) then normalise, and return the Rayleigh quotient v . hvp(v) as a float.

Function signature

def hvp(f, x, v):
def top_hessian_eigenvalue(f, x, iters):

Example

Input

A = torch.tensor([[2.0, 1.0], [1.0, 3.0]])
f = lambda x: 0.5 * x @ A @ x
hvp(f, torch.tensor([1.0, -1.0]), torch.tensor([1.0, 0.0]))

Output

tensor([2., 1.])

Reasoning

The Hessian of 0.5 * x.T A x is A for a symmetric A, at every x. So H @ v is the first column of A.

Constraints and notes

• f is built from torch operations and returns a scalar tensor. x and v have shape (n,), and x may not require a gradient.
• Return a detached tensor of shape (n,). Do not change x or leave a gradient on it.
• n can be large (hundreds of thousands), so never build H.