laplax - Laplace Approximations with JAX

Tobias Weber, Bálint Mucsányi, Lenard Rommel, Thomas Christie, Lars Kasüschke, Marvin Pförtner, Philipp Hennig

Keywords: Laplace Approximation, JAX, Uncertainty Quantification, Bayesian Deep Learning, Software

TL;DR: Software library 'laplax' for Laplace Approximations with JAX

Abstract: The Laplace approximation provides a scalable and efficient means of quantifying weight-space uncertainty in deep neural networks, enabling the application of Bayesian tools such as predictive uncertainty and model selection via Occam's razor. In this work, we introduce [`laplax`](https://anonymous.4open.science/r/laplax_anonymous-0472/), a new open-source Python package for performing Laplace approximations in `jax`. Designed with a modular and purely functional architecture and minimal external dependencies, `laplax` offers a flexible and researcher-friendly framework for rapid prototyping and experimentation. Its goal is to facilitate research on Bayesian neural networks, uncertainty quantification for deep learning, and the development of improved Laplace approximation techniques.

Submission Number: 35