laplax -- Laplace Approximations with JAX

2025-07-22

AI Summary

This work addresses the inefficiency and poor scalability of weight uncertainty quantification in Bayesian deep learning. We propose an efficient Laplace approximation framework built on JAX, leveraging its automatic differentiation and functional programming capabilities to enable modular, low-dependency second-order posterior approximation. Our method supports rapid computation of posterior weight covariance for large-scale neural networks, enabling principled prediction uncertainty estimation and model selection via the Occam’s razor principle. Our primary contribution is laplax—an open-source, lightweight toolkit that systematically adapts the Laplace approximation to modern deep learning architectures and distributed training paradigms. By preserving theoretical rigor while drastically reducing implementation complexity, laplax lowers the barrier to Bayesian inference and facilitates practical uncertainty modeling and algorithmic advancement in large language models and other foundation models.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Probabilistic Programming
• Natural Language Processing: (Large) Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Web Mining and Content Analysis: Large pretrained models with web data

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, a new open-source Python package for performing Laplace approximations with 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.

Problem

• Quantify weight-space uncertainty in deep neural networks
• Provide scalable Bayesian tools for predictive uncertainty
• Facilitate research on improved Laplace approximation techniques

Innovation

• Laplace approximation for uncertainty quantification
• Modular functional architecture with JAX
• Open-source Python package for rapid prototyping