JaxSGMC: Modular stochastic gradient MCMC in JAX
Not on the map. Why?
Abstract (arXiv)
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.
Key facts. Cited by 1 paper in the map (data through 14 Oct 2026). Builds on 29 of its 49 references that are in the map. First posted on arXiv 16 May 2025.
Builds on
29 of 49 references.
References · 29
- 2015Deep Residual Learning for Image Recognitioncited by 52,096
- 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,399
- 2015Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learningcited by 4,448
- 2016Simple and Scalable Predictive Uncertainty Estimation using Deep Ensemblescited by 3,926
- 2017MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applicationscited by 3,845
- 2016Distribution-Free Predictive Inference For Regressioncited by 668
- 2020A Survey of Deep Active Learningcited by 554
- 2020Bayesian Deep Learning and a Probabilistic Perspective of Generalizationcited by 480
- 2019A Survey of Deep Learning Techniques for Autonomous Drivingcited by 403
- 2018Pyro: Deep Universal Probabilistic Programmingcited by 380
- 2014Stochastic Gradient Hamiltonian Monte Carlocited by 343
- 2019Cyclical Stochastic Gradient MCMC for Bayesian Deep Learningcited by 203
17 more
- 2019Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyrocited by 191
- 2015A Complete Recipe for Stochastic Gradient MCMCcited by 146
- 2015Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural Networkscited by 121
- 2019Machine learning for molecular simulationcited by 110
- 2017TensorFlow Distributionscited by 106
- 2014Consistency and fluctuations for stochastic gradient Langevin dynamicscited by 65
- 2019Stochastic gradient Markov chain Monte Carlocited by 56
- 2020Non-convex Learning via Replica Exchange Stochastic Gradient MCMCcited by 34
- 2022NeuralUQ: A comprehensive library for uncertainty quantification in neural differential equations and operatorscited by 33
- 2021Exact Langevin Dynamics with Stochastic Gradientscited by 30
- 2021Learning neural network potentials from experimental data via Differentiable Trajectory Reweightingcited by 26
- 2018Stochastic Gradient MCMC with Repulsive Forcescited by 18
- 2023Exploring Distributional Shifts in Large Language Models for Code Analysiscited by 14
- 2020AMAGOLD: Amortized Metropolis Adjustment for Efficient Stochastic Gradient MCMCcited by 13
- 2022Scalable Bayesian Uncertainty Quantification for Neural Network Potentials: Promise and Pitfallscited by 11
- 2020Bayesian Graph Neural Networks for Molecular Property Predictioncited by 8
- 2017sgmcmc: An R Package for Stochastic Gradient Markov Chain Monte Carlocited by 3
Built on by
Citing papers per year: 2026: 1.
1 paper.