ICML 2025

Efficiently Vectorized MCMC on Modern Accelerators
Hugh Dance 1 Pierre Glaser 1 Peter Orbanz 1 Ryan Adams 2
Abstract
With the advent of automatic vectorization tools
(e.g., JAX’s vmap), writing multi-chain MCMC
algorithms is often now as simple as invoking
those tools on single-chain code. Whilst conve-
nient, 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 vectoriz-
ing with tools like vmap, by using the framework
of finite state machines (FSMs). Using a simpli-
fied 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 De-
layed Rejection, demonstrating speed-ups of up
to an order of magnitude in experiments.
1. Introduction
Automatic vectorization is the act of transforming a func-
tion to handle batches of inputs without user intervention.
Implementations of automatic vectorization algorithms are
now available in many mainstream scientific computing li-
braries, and have dramatically simplified the task of running
multiple instances of a single algorithm concurrently. They
are routinely used to train neural networks (Flax, 2023) and
in other scientific applications, e.g., Schoenholz and Cubuk
(2021); Oktay et al. (2023); Pfau et al. (2020).
This paper focuses on the use of automatic vectorization
for Markov chain Monte Carlo (MCMC) algorithms. Tools
like JAX’s vmap provide a convenient way to run multiple
1Gatsby Unit, University College London, UK 2Department of
Computer Science, Princeton University, USA. Correspondence
to: Hugh Dance <hugh.dance.15@ucl.ac.uk>.
Proceedings of the 4...