JaxSGMC: Modular stochastic gradient MCMC in JAX

2024-05-01

SoftwareX

AI Summary

To address the scalability limitations of uncertainty quantification (UQ) in Bayesian deep learning for large-scale, high-dimensional settings, this work introduces the first modular stochastic gradient Markov chain Monte Carlo (SG-MCMC) framework designed natively for JAX. Our method features a highly decoupled architecture that enables plug-and-play composition and automatic differentiation compatibility across core components—including gradient noise models, preconditioners, and numerical integrators. Leveraging JAX’s functional programming paradigm, we achieve efficient cross-device (CPU/GPU/TPU) parallel sampling via just-in-time compilation (jit), vectorization (vmap), automatic differentiation (grad), and multi-process mapping (pmap). Empirical evaluation on multiple Bayesian neural network benchmarks demonstrates a 3–5× speedup in sampling throughput and a 40% reduction in memory consumption compared to PyTorch- and TensorFlow-based SG-MCMC implementations, substantially lowering the barrier to deploying SG-MCMC in trustworthy deep learning systems.

Technology Category

Machine Learning: Calibration & Uncertainty QuantificationReasoning under Uncertainty: Stochastic OptimizationSearch and Optimization: Sampling/Simulation-based Search

Application Category

User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendationGraph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphsSearch and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

Problem

• Develops JaxSGMC for scalable Bayesian deep learning
• Reduces barriers to adopt SG-MCMC sampling methods
• Enables modular custom sampler design for SG-MCMC

Innovation

• Modular SG-MCMC library in JAX
• State-of-the-art SG-MCMC samplers
• Custom samplers from standard blocks