Hessian calculation

[stefano-mossa]

I have two (probably) trivial questions:

1. is it possible to compute one element only of the hessian (h_ij) avoiding to calculate the full Hessian matrix with jacfwd(jacrev(fun))?
2. In any case, jacfwd(jacrev(fun)) is not exploiting any symmetry of the hessian in the computation, correct?

[mattjj]

Great questions!

For your first question, the answer is yes, basically by using autodiff to get access to the function (u, v) -> u'Hv. With that function you can extract any element of the Hessian by feeding in standard basis vectors for u and v.

More concretely:

import jax
import jax.numpy as jnp

def hessian_bilinear_form(f, x, u, v):
  # forward-over-reverse autodiff
  f_jvp = lambda x: jax.jvp(f, (x,), (u,))[1]
  return jax.jvp(f_jvp, (x,), (v,))[1]

def hessian_entry(f, x, i, j):
  u = jnp.zeros_like(x).at[i].set(1)
  v = jnp.zeros_like(x).at[j].set(1)
  return hessian_bilinear_form(f, x, u, v)

A = jnp.array([[1., 2],
               [3., 4.]])
f = lambda x: jnp.dot(x, jnp.dot(A, x))

H_01 = hessian_entry(f, jnp.array([1., 2.]), 0, 1)
print(H_01)

H = A + A.T
print(H[0, 1])

For your second question, that's correct: it's not exploiting any symmetry. The jet function in jax.experimental.jet does exploit symmetry, which becomes important at high order autodiff; see the discussion in #3730. I'm not sure if it'd help much at second order though.

Let me know if that covers what you had in mind, or if you spot any mistakes!

[stefano-mossa]

It works perfectly indeed. For the record, I am trying to workaround a sparse representation of a (very sparse) hessian without going through the dense matrix. Thank you very much for your help.