<p>Not sure if this will help in your particular use-case, but you could work with an approximation of the Hessian, e.g. empirical Fisher (EF). I've worked with this approach to implement Laplace approximation for Flux models (see <a href="https://github.com/pat-alt/BayesLaplace.jl" rel="nofollow noreferrer">here</a>) inspired by this <a href="https://github.com/AlexImmer/Laplace" rel="nofollow noreferrer">PyTorch implementation</a>. Below I've applied the approach to your example.</p>
<pre><code>using Flux: Chain, Dense, σ, crossentropy, params, DataLoader
using Zygote
using Random

Random.seed!(2022)
model = Chain(
    x -&gt; reshape(x, :, size(x, 4)),
    Dense(2, 5),
    Dense(5, 1),
    x -&gt; σ.(x)
)
n_data = 5
input = randn(2, 1, 1, n_data)
target = randn(1, n_data)
loss(x, y) = crossentropy(model(x), y)

n_params = length(reduce(vcat, [vec(θ) for θ ∈ params(model)]))
 = zeros(n_params,n_params)
data = DataLoader((input, target))

for d in data
  x, y = d
   = gradient(() -&gt; loss(x,y),params(model))  
   = reduce(vcat,[vec([θ]) for θ ∈ params(model)])
   +=  * ' # empirical fisher
end
</code></pre>
<p>Should there be a way to use Zygote autodiff directly (and more efficiently) I'd also be interested to see that. Using EF for the full Hessian still scales quadratically in the number of parameters, but as shown in this NeurIPS 2021 <a href="https://arxiv.org/abs/2106.14806" rel="nofollow noreferrer">paper</a> you can further approximate the Hessian using (blog-)diagonal factorization. The paper also shows that in the context of Bayesian deep learning treating only the last layer probabilistically generally yields good results, but again not sure if relevant in your case.</p>
Calculating Hessian of a loss function including a NN w.r.t parameters in Julia using Zygote
<p>How would you calculate a hessian of a loss function that consists of a Neural Network w.r.t. the NN's parameters?</p>
<p>For instance, consider the loss function below</p>
<pre><code>using Flux: Chain, Dense, σ, crossentropy, params
using Zygote
model = Chain(
    x -&gt; reshape(x, :, size(x, 4)),
    Dense(2, 5),
    Dense(5, 1),
    x -&gt; σ.(x)
)
n_data = 5
input = randn(2, 1, 1, n_data)
target = randn(1, n_data)
loss = model -&gt; crossentropy(model(input), target)
</code></pre>
<p>I can get a gradient w.r.t parameters in two ways…</p>
<pre><code>Zygote.gradient(model -&gt; loss(model), model)
</code></pre>
<p>or</p>
<pre><code>grad = Zygote.gradient(() -&gt; loss(model), params(model))
grad[params(model)[1]]
</code></pre>
<p>However, I can’t find a way to get a hessian w.r.t its parameters. (I want to do something like <code>Zygote.hessian(model -&gt; loss(model), model)</code>, but <code>Zygote.hessian</code> does not take <code>::Params</code> as an input)</p>
<p>Recently, a <code>jacobian</code> function <a href="https://github.com/FluxML/Zygote.jl/issues/910" rel="nofollow noreferrer">was added</a> to the master branch (issue #910), which <a href="https://github.com/FluxML/Zygote.jl/pull/414" rel="nofollow noreferrer">understands <code>::Params</code> as an input</a>.</p>
<p>I've been trying to combine <code>gradient</code> and <code>jacobian</code> to get a hessian (because a hessian is the jacobian of a gradient of a function), but to no avail.
I think the problem is that <code>model</code> is a <code>Chain</code> object that includes generic functions like <code>reshape</code> and <code>σ.</code> which lack parameters, but I can't get past this.</p>
<pre><code>grad = model -&gt; Zygote.gradient(model -&gt; loss(model), model)
jacob = model -&gt; Zygote.jacobian(grad, model)
jacob(model) ## does not work
</code></pre>
<p>EDIT: For reference, I've <a href="https://github.com/stash-196/pytorch-derivatives/blob/main/derivatives.py" rel="nofollow noreferrer">created this in pytorch before</a></p>