JaxSGMC: Modular stochastic gradient MCMC in JAX
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2025
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866915289922273280 |
|---|---|
| author | Thaler, Stephan Fuchs, Paul Cukarska, Ana Zavadlav, Julija |
| author_facet | Thaler, Stephan Fuchs, Paul Cukarska, Ana Zavadlav, Julija |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2505_11190 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | JaxSGMC: Modular stochastic gradient MCMC in JAX Thaler, Stephan Fuchs, Paul Cukarska, Ana Zavadlav, Julija Computation Applications Machine Learning 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. |
| title | JaxSGMC: Modular stochastic gradient MCMC in JAX |
| topic | Computation Applications Machine Learning |
| url | https://arxiv.org/abs/2505.11190 |