Can this be used for Laplace approximation?

[blackblitz]

In Laplace approximation, the Hessian of the loss function is computed for quadratic approximation. Can this package be used to do a block-diagonal approximation of the Hessian at the minimum? If yes, could you please show (using jax and flax) how to approximate it and define a quadratic approximation of the loss function (which should be something like 1/2 (theta - theta_star)^T H(L)(theta_star) (theta - theta_star), where theta_star is the minimum and H(L) is the Hessian of the loss function)?

[botev]

Hi, yes, this can be used for a Laplace approximation. In particular you can take a look at the CurvatureEstimator and how it can be used. We don't have code for that, but it can be fairly straightforward to do it. There are however several details that one might need to pay attention to, such as:

1. The code in the library estimates the "average" GGN/Fisher, which means that if you have a prior, you need to rescale it by 1/N
2. The approximation to the GGN might not be a perfect approximation to the full GGN, so you might want to consider techniques which can mitigate that.