Hessian-vector products

[baggepinnen]

This is more or less a copy of FluxML/Flux.jl#129, adopted to Zygote in order to track the progress of HVPs

using Zygote: gradient
inp     = randn(3)             # Input
v       = randn(3)             # Vector
H       = randn(3,3); H = H+H' # Hessian
f(inp,H)  = 0.5*inp'H*inp      # i'H*i function to take hessian of
hvp     = H*v                  # True Hessian vector product
gg      = H*inp                # True gradient
ggvp    = gg'v                 # True gradient vector product
Zygote.refresh()
g(x)    = gradient(x->f(x,H),x)[1]  # gradient function
gvp(x)  = g(x)⋅v
@assert gvp(inp) ≈ ggvp        # Correct until here == gg'v
dgvp(x) = gradient(gvp, x)
dgvp(inp)                      # Errors

Can't differentiate foreigncall expression
error(::String) at error.jl:33
get at abstractdict.jl:595 [inlined]
(::typeof(∂(get)))(::Nothing) at interface2.jl:0
in at abstractdict.jl:665 [inlined]
(::typeof(∂(in)))(::Nothing) at interface2.jl:0
haskey at abstractdict.jl:17 [inlined]
(::typeof(∂(haskey)))(::Nothing) at interface2.jl:0
macro expansion at lib.jl:37 [inlined]
accum_param at lib.jl:35 [inlined]
(::typeof(∂(Zygote.accum_param)))(::Nothing) at interface2.jl:0
#137 at lib.jl:58 [inlined]
(::typeof(∂(λ)))(::Nothing) at interface2.jl:0
#209#back at grad.jl:41 [inlined]
(::typeof(∂(λ)))(::Nothing) at interface2.jl:0
#19 at hvp.jl:10 [inlined]
(::typeof(∂(λ)))(::Tuple{Nothing,Array{Float64,1}}) at interface2.jl:0
#73 at interface.jl:38 [inlined]
(::typeof(∂(λ)))(::Tuple{Array{Float64,1}}) at interface2.jl:0
gradient at interface.jl:44 [inlined]
(::typeof(∂(Zygote.gradient)))(::Tuple{Array{Float64,1}}) at interface2.jl:0
g at hvp.jl:10 [inlined]
(::typeof(∂(g)))(::Array{Float64,1}) at interface2.jl:0
gvp at hvp.jl:11 [inlined]
(::typeof(∂(gvp)))(::Int8) at interface2.jl:0
(::getfield(Zygote, Symbol("##73#74")){typeof(∂(gvp))})(::Int8) at interface.jl:38
gradient(::Function, ::Array{Float64,1}) at interface.jl:44
dgvp(::Array{Float64,1}) at hvp.jl:13
top-level scope at none:0

[MikeInnes]

How about doing this with ForwardDiff over Zygote? Should be doable in a similar way to Hessians now.

[baggepinnen]

Yes it now works with ForwardDiff on Zygote#master

] add Zygote#master
using Zygote
using Zygote: gradient, forward_jacobian
inp     = randn(3)             # Input
v       = randn(3)             # Vector
H       = randn(3,3); H = H+H' # Hessian
f(inp,H)  = 0.5*inp'H*inp      # i'H*i function to take hessian of
hvp     = H*v                  # True Hessian vector product
gg      = H*inp                # True gradient
ggvp    = gg'v                 # True gradient vector product
Zygote.refresh()
g(x)    = gradient(x->f(x,H),x)[1]  # gradient function
gvp(x)  = g(x)⋅v
@assert gvp(inp) ≈ ggvp        # Correct until here == gg'v
dgvp(x) = forward_jacobian(gvp, x)[2]
@assert dgvp(inp) ≈ hvp        # Works

[baggepinnen]

Want a PR with the following function?

function hessian_vector_product(f,v,x)
    g    = x -> gradient(f,x)[1]  # gradient function
    gvp  = x -> g(x)⋅v
    forward_jacobian(gvp, x)[2]
end

@assert hessian_vector_product(f,v,inp) ≈ hvp # true

Alternatively, a function hvp that returns a function to evaluate the hvp?

Edit: Actually, it seems there is no performance benefit to my suggested function over just calling hessian(f,inp)*v, not even if hvpf = hvp(f,v) is precalculated and applied as hvpf(inp)

[MikeInnes]

I think you can do an HVP via custom seeding of the dual numbers, which should be faster, so that might be worth exploring.

[ChrisRackauckas]

Note that the example from #250 is from https://github.com/JuliaDiffEq/SparseDiffTools.jl which is implementing Hessian-vector products and ran into some issues, and so we are tracking an implementation there.

[andreasko]

Yes it now works with ForwardDiff on Zygote#master

] add Zygote#master
using Zygote
using Zygote: gradient, forward_jacobian
inp     = randn(3)             # Input
v       = randn(3)             # Vector
H       = randn(3,3); H = H+H' # Hessian
f(inp,H)  = 0.5*inp'H*inp      # i'H*i function to take hessian of
hvp     = H*v                  # True Hessian vector product
gg      = H*inp                # True gradient
ggvp    = gg'v                 # True gradient vector product
Zygote.refresh()
g(x)    = gradient(x->f(x,H),x)[1]  # gradient function
gvp(x)  = g(x)⋅v
@assert gvp(inp) ≈ ggvp        # Correct until here == gg'v
dgvp(x) = forward_jacobian(gvp, x)[2]
@assert dgvp(inp) ≈ hvp        # Works

Hey all,
I tried the ForwardDiff over Zygote example but ran into an error:
The function call g(x) = gradient(x->f(x,H),x) breaks with

"Output should be scalar; gradients are not defined for output [...]"

Here, [...] is a vector with one element.

[DhairyaLGandhi]

You might want to add a reduction to the output of the function

[tomoleary]

If this is of any use: the following seems to work for Hessian-matrix products using Zygote

using Zygote
n = 100
r = 3
x = randn(n)
f(x) = prod([cos(x[i]) for i = 1 : n])
g(x) = Zygote.gradient(x->f(x),x)[1]
H(x) = Zygote.hessian(x->f(x),x)
v = rand(n,r)
Hvp(x,v) = Zygote.forward_jacobian(x->g(x)'v,x)[2]
error = Hvp(x,v) - H(x)*v
println("||error||_2 = ",opnorm(error))