# JaxSGMC: Modular stochastic gradient MCMC in JAX

arXiv: https://arxiv.org/abs/2505.11190 · HTML: https://mlsanity.com/paper/2505.11190/

## Key facts

- Cited by 1 papers 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 (most cited first, 29 of 29)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications](https://mlsanity.com/paper/1704.04861/) (2017)
- [Distribution-Free Predictive Inference For Regression](https://mlsanity.com/paper/1604.04173/) (2016)
- [A Survey of Deep Active Learning](https://mlsanity.com/paper/2009.00236/) (2020)
- [Bayesian Deep Learning and a Probabilistic Perspective of Generalization](https://mlsanity.com/paper/2002.08791/) (2020)
- [A Survey of Deep Learning Techniques for Autonomous Driving](https://mlsanity.com/paper/1910.07738/) (2019)
- [Pyro: Deep Universal Probabilistic Programming](https://mlsanity.com/paper/1810.09538/) (2018)
- [Stochastic Gradient Hamiltonian Monte Carlo](https://mlsanity.com/paper/1402.4102/) (2014)
- [Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning](https://mlsanity.com/paper/1902.03932/) (2019)
- [Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro](https://mlsanity.com/paper/1912.11554/) (2019)
- [A Complete Recipe for Stochastic Gradient MCMC](https://mlsanity.com/paper/1506.04696/) (2015)
- [Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural Networks](https://mlsanity.com/paper/1512.07666/) (2015)
- [Machine learning for molecular simulation](https://mlsanity.com/paper/1911.02792/) (2019)
- [TensorFlow Distributions](https://mlsanity.com/paper/1711.10604/) (2017)
- [Consistency and fluctuations for stochastic gradient Langevin dynamics](https://mlsanity.com/paper/1409.0578/) (2014)
- [Stochastic gradient Markov chain Monte Carlo](https://mlsanity.com/paper/1907.06986/) (2019)
- [Non-convex Learning via Replica Exchange Stochastic Gradient MCMC](https://mlsanity.com/paper/2008.05367/) (2020)
- [NeuralUQ: A comprehensive library for uncertainty quantification in neural differential equations and operators](https://mlsanity.com/paper/2208.11866/) (2022)
- [Exact Langevin Dynamics with Stochastic Gradients](https://mlsanity.com/paper/2102.01691/) (2021)
- [Learning neural network potentials from experimental data via Differentiable Trajectory Reweighting](https://mlsanity.com/paper/2106.01138/) (2021)
- [Stochastic Gradient MCMC with Repulsive Forces](https://mlsanity.com/paper/1812.00071/) (2018)
- [Exploring Distributional Shifts in Large Language Models for Code Analysis](https://mlsanity.com/paper/2303.09128/) (2023)
- [AMAGOLD: Amortized Metropolis Adjustment for Efficient Stochastic Gradient MCMC](https://mlsanity.com/paper/2003.00193/) (2020)
- [Scalable Bayesian Uncertainty Quantification for Neural Network Potentials: Promise and Pitfalls](https://mlsanity.com/paper/2212.07959/) (2022)
- [Bayesian Graph Neural Networks for Molecular Property Prediction](https://mlsanity.com/paper/2012.02089/) (2020)
- [sgmcmc: An R Package for Stochastic Gradient Markov Chain Monte Carlo](https://mlsanity.com/paper/1710.00578/) (2017)

## Built on by (most cited first, 1 of 1)

- [Position: The Time for Sampling Is Now! Charting a New Course for Bayesian Deep Learning](https://mlsanity.com/paper/2605.21765/) (2026)

## 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.