Laplace approximation

[theorashid]

Part of INLA roadmap #340.

This would be a Laplace approximation for the entire posterior. Perhaps the API can be like the current PyMC VI API.

Notes

• Previous PyMC attempt Adds quadratic approximation pymc#4847
• pymc3 effort
• jax implementation
• Dan's blog

[carsten-j]

I am new to PYMC and probabilistic programming. While reading Statistical Rethinking I was looking for a Quadratic Approximation method for Python. I did not find any. Partly based on the first bullet point above I have created this package https://github.com/carsten-j/quad5. I have based in on the PYMC class ArrayStep such that I can use the sample method on PYMC model and automatically get constant, deterministic, and observed data in the inferencedata output from the sample method. Feedback is appreciated.

[theorashid]

Hi Carsten,

Sorry I was slow on this – I was on holiday. This is really cool.

Reading through the README, yes the AutoLaplaceApproximation from numpyro is slightly different. It uses quadratic on all the parameters (I think), in more of a variational inference set up. The main numpyro dev ported the examples from Statistical Rethinking and used that for the relevant chapters, e.g. here. The equivalent variational inference method for pymc would lie in the pm.fit() API, something like pm.fit(method="laplace") or pm.fit(method="quadratic"). I would still like to see that in pymc at some point.

I was unfamiliar with ArrayStep, but this fits nicely. I'm not sure if it's possible in this set up, but it would be great if ArrayStep allows us to do quadratic approximation on a subset of the parameters. This would actually solve the #344 (marginal Laplace approx). Also I didn't know much about pm.find_MAP() so if that's all in pytensor, it might also cover #342 (Newton solver to find mode) because it is all differentiable.

If you're willing, we would love you to contribute this to pymc-experimental. In this way, we can more formally review the code with the setup here (@elizavetasemenova, @zaxtax and I can review). There are step-by-step instructions on contributing to pymc here, but if there's anything you're struggling with I'm happy to help.

Really excited to get this in.

[carsten-j]

Hi. Thanks for the positive feedback. I am a bit busy for the next couple of days but after Wednesday I will start looking into how to contribute to PYMC experimental.

[theorashid]

Excellent, looking forward to it

[ricardoV94]

@carsten-j if we want to use Laplace as a step, it should probably fit in every iteration, specially if the step sampler is assigned only a subset of the model variables, since it should be conditioned on the rest of the variables. Also with stuff like pm.Minibatch the logp won't be constant across iterations, even if all the variables are assigned to the Laplace step sampler.

As mentioned by @theorashid it probably makes more sense to have a pm.fit(method="laplace") first. The step sampler idea may be useful specifically for partial approximations (a subset of variables), but for a simple approximation it's a bit of an overkill.

Obtaining the the rest of the InferenceData metadata shouldn't be too hard, if you dig through the codebase you'll see that that happens after sampling anyway with pm.to_inferencedata: https://github.com/pymc-devs/pymc/blob/e5c3030d99136a23d910a4b2308e5750c7937aca/pymc/sampling/mcmc.py#L941

[carsten-j]

I am using a Mac and Visual Studio Code. I can run tests from the commandline and I can also run tests from within Visual Studio Code. But when I add a breakpoint to a test and try to debug the test from Visual Studio Code then I get this error message

message = MatplotlibDeprecationWarning('The interactive_bk attribute was deprecated in Matplotlib 3.9 and will be removed in 3.11. Use ``matplotlib.backends.backend_registry.list_builtin(matplotlib.backends.BackendFilter.INTERACTIVE)`` instead.')

Do you know if debugging tests from Visual Studio Code should work?

[theorashid]

Not sure on your specific case, but I use VSCode's interface to run specific tests. I know VSCode has a test debug interface too which I haven't used but the comments here might help https://stackoverflow.com/questions/63619498/vs-code-python-debugging-pytest-test-with-the-debugger