Forward over reverse off-diagonal Hessian-vector product

[qsh-zh]

I have a scalar function

I am interested in the following term

My questions as follows:

1. "forward-over-reverse" off-diagonal Hessian-vector product. According to the tutorial, we should have an efficient approach based on forward-over-reverse, but how to do it for this off -digonal Hessian matrix,
2. The vector part and hessian are highly corrected, can we save some calculation? For example, save get gradients when calculating the hessian term?

[YouJiacheng]

1. See below.
2. Yes, jvp over grad just do this. But for hessian-grad product, we need to use linearize.

def f(y, x):
    return jnp.sum(y) + jnp.sum(x) # just example

def g(y, x, u):
    f_x_func = jax.grad(f, argnums=1) # autodiff for x only
    f_x, f_yx_dot_u = jax.jvp(lambda _y: f_x_func(_y, x), (y,), u)  # use closure to autodiff for y only
    # \frac{\partial f(y,x)}}{\partial x} and \frac{\partial^2 f(y,x)}}{\partial y\partial x} u

def h(y, x, u):
    f_x_func = jax.grad(f, argnums=1) # autodiff for x only
    f_x, f_yx_func = jax.linearize(lambda _y: f_x_func(_y, x), y)  # use closure to autodiff for y only
    return f_yx_func(f_x - u)

[qsh-zh]

@YouJiacheng Thanks for the clear explanation.

[qsh-zh]

@YouJiacheng Sorry, I think there is a small mistake in your answer. In the last line, f_yx_func can work on primal value, whose shape should similar to the y.shape.

[YouJiacheng]

@YouJiacheng Sorry, I think there is a small mistake in your answer. In the last line, f_yx_func can on work on primal value, whos shape should similar to the y.shape.

Oh, I just rephrase your formula and not see that \frac{\partial^2 f(y,x)}}{\partial y\partial x} has shape m×n.
I think if \frac{\partial^2 f(y,x)}}{\partial y\partial x} is \frac{\partial}{\partial y}\frac{\partial f(y,x)}}{\partial x} it can not dot with tangent with shape==x.shape.
I think you can change the order of differentiation, or use

def h(y, x, u):
    f_x_func = jax.grad(f, argnums=1) # autodiff for x only
    f_x, f_yx_func = jax.linearize(lambda _y: f_x_func(_y, x), y)  # use closure to autodiff for y only
    return jax.linear_transpose(f_yx_func, y)(f_x - u) # linear_transpose only use y.shape and y.dtype

[qsh-zh]

But actually, if we use linear_transpose, it is equivalent to reverse over the reverse. Seems forward over reverse is not a good option here.

[YouJiacheng]

But actually, if we use linear_transpose, it is equivalent to reverse over the reverse. Seems forward over reverse is not a good option here.

I agree that linearize + linear_transpose is equivalent to vjp(but linear_transpose itself isn't). Actually, linearize costs about same memory as reverse-mode AD.

I think it is efficient to change the partial differentiation order(assume C^2 smooth), and sequentially evaluate the gradient value using a separate call to grad function and the hessian-vector product value using jvp(abort the out primal). The redundancy is caused by the sequential dependence between hessian-vector product and evaluate gradient since your vector contain the gradient.