# Efficiently Vectorized MCMC on Modern Accelerators

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 14 of its 40 references that are in the map.
- First posted on arXiv 20 Mar 2025.

## Builds on (most cited first, 14 of 14)

- [TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems](https://mlsanity.com/paper/1603.04467/) (2016)
- [Ab-Initio Solution of the Many-Electron Schr"odinger Equation with Deep Neural Networks](https://mlsanity.com/paper/1909.02487/) (2019)
- [Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro](https://mlsanity.com/paper/1912.11554/) (2019)
- [NeuTra-lizing Bad Geometry in Hamiltonian Monte Carlo Using Neural Transport](https://mlsanity.com/paper/1903.03704/) (2019)
- [Learning Finite State Representations of Recurrent Policy Networks](https://mlsanity.com/paper/1811.12530/) (2018)
- [BlackJAX: Composable Bayesian inference in JAX](https://mlsanity.com/paper/2402.10797/) (2024)
- [Automatic Reparameterisation of Probabilistic Programs](https://mlsanity.com/paper/1906.03028/) (2019)
- [Transport Elliptical Slice Sampling](https://mlsanity.com/paper/2210.10644/) (2022)
- [tfp.mcmc: Modern Markov Chain Monte Carlo Tools Built for Modern Hardware](https://mlsanity.com/paper/2002.01184/) (2020)
- [Delayed rejection Hamiltonian Monte Carlo for sampling multiscale distributions](https://mlsanity.com/paper/2110.00610/) (2021)
- [Automatically Batching Control-Intensive Programs for Modern Accelerators](https://mlsanity.com/paper/1910.11141/) (2019)
- [Recurrent Neural Language Models as Probabilistic Finite-state Automata](https://mlsanity.com/paper/2310.05161/) (2023)
- [Infinite Factorial Finite State Machine for Blind Multiuser Channel Estimation](https://mlsanity.com/paper/1810.09261/) (2018)
- [Neuromechanical Autoencoders: Learning to Couple Elastic and Neural Network Nonlinearity](https://mlsanity.com/paper/2302.00032/) (2023)

## Abstract (arXiv)

With the advent of automatic vectorization tools (e.g., JAX's textttvmap), writing multi-chain MCMC algorithms is often now as simple as invoking those tools on single-chain code. Whilst convenient, for various MCMC algorithms this results in a synchronization problem -- loosely speaking, at each iteration all chains running in parallel must wait until the last chain has finished drawing its sample. In this work, we show how to design single-chain MCMC algorithms in a way that avoids synchronization overheads when vectorizing with tools like textttvmap by using the framework of finite state machines (FSMs). Using a simplified model, we derive an exact theoretical form of the obtainable speed-ups using our approach, and use it to make principled recommendations for optimal algorithm design. We implement several popular MCMC algorithms as FSMs, including Elliptical Slice Sampling, HMC-NUTS, and Delayed Rejection, demonstrating speed-ups of up to an order of magnitude in experiments.