Avoiding redundant computation when using Zygote and ForwardDiff for forward-over-reverse hessians

IvanBioli
Tue, 11 Mar 2025 19:15:11 +0000
General Usage

I am computing the Hessian of a function using a combination of Zygote (reverse-mode AD) and ForwardDiff (forward-mode AD). A MWE of my current implementation is as follows:

using Zygote, ForwardDiff, DiffResults

# Define a simple function that takes a vector and returns a scalar
f_aux(x) = sum(x .^ 2)  # Example: sum of squares

# Gradient function
function gfun(x)
    valgrad = Zygote.withgradient(f_aux, x) # Reverse-mode autodiff
    return valgrad.grad[1]
end

# Hessian function
hessfun = x -> ForwardDiff.jacobian!(DiffResults.JacobianResult(x), gfun, x)  # Forward-mode autodiff (forward-over-reverse)

# Example input vector
x = rand(2)
hessfun(x)

This implementation ensures that I don’t recompute the gradient when computing the Hessian. However, I still have to recompute f_aux(x) separately if I need its value along with the Hessian.

Is there a way to modify this approach so that I can efficiently retrieve both f_aux(x) and the Hessian without recomputing f_aux(x)? Ideally, I would like a solution that remains compatible with Zygote and ForwardDiff, as these are the only AD tools I have found to work with Lux.jl and CUDA.jl for GPU acceleration. In my use case, f_aux(x) is actually a neural network (NN) where x are the NN’s input, and they are low-dimensional, and the NN output is scalar.

Any insights or suggestions would be greatly appreciated!

5 posts - 2 participants