JaxSGMC: Modular stochastic gradient MCMC in JAX

Abstract

We present JaxSGMC, an application-agnostic library for stochastic gradient Markov chain Monte Carlo (SG-MCMC) in JAX. SG-MCMC schemes are uncertainty quantification (UQ) methods that scale to large datasets and high-dimensional models, enabling trustworthy neural network predictions via Bayesian deep learning. JaxSGMC implements several state-of-the-art SG-MCMC samplers to promote UQ in deep learning by reducing the barriers of entry for switching from stochastic optimization to SG-MCMC sampling. Additionally, JaxSGMC allows users to build custom samplers from standard SG-MCMC building blocks. Due to this modular structure, we anticipate that JaxSGMC will accelerate research into novel SG-MCMC schemes and facilitate their application across a broad range of domains.

References in corpus (9)

• MCMC using Hamiltonian dynamics
• A Survey of Deep Learning Techniques for Autonomous Driving
• Machine learning for molecular simulation
• Pyro: Deep Universal Probabilistic Programming
• TensorFlow Distributions
• A Complete Recipe for Stochastic Gradient MCMC
• Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning
• Deep Coarse-grained Potentials via Relative Entropy Minimization
• NeuralUQ: A comprehensive library for uncertainty quantification in neural differential equations and operators

Cited by in corpus (1)

• chemtrain: Learning Deep Potential Models via Automatic Differentiation and Statistical Physics