Research map / Learning theory and algorithms / Prediction and decision losses
Efficiently Vectorized MCMC on Modern Accelerators
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Converts single-chain MCMC samplers with while loops into finite state machines so that vmap-vectorized chains can avoid per-iteration synchronization barriers, giving up to ~10x speed-ups.
Problem
Automatic vectorization (e.g., JAX vmap) of multi-chain MCMC algorithms with variable-length while loops forces all chains to wait for the slowest chain at every iteration, wasting accelerator capacity.
Approach
Decompose each MCMC transition kernel into contiguous code blocks executed as states of a finite state machine, with a step function that lets each chain advance through its own block sequence so only the n-sample collection requires synchronization; adds step bundling and cost amortization to reduce overhead, plus a theoretical speed-up bound R(m)=E max N / E N.
Key results
- FSM elliptical slice sampler runs ~3x faster than BlackJAX at 1024 chains on Real Estate GP, reducing 10k-sample time from over half an hour to about 10 minutes.
- Condensed FSM for symmetric delayed rejection achieves near order-of-magnitude speed-ups at m=1024 chains, with bundling giving ~3x efficiency gain over unbundled FSM.
- FSM NUTS gives nearly an order of magnitude ESS/sec speed-up for 100 chains on a 100-dimensional correlated mixture, and about half an order of magnitude for 500 chains.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use an FSM implementation instead of standard NUTS | NUTS · Predator Prey (PP) | speed-up: 1.5× | Helped |
| Use an FSM implementation instead of standard TESS | TESS · Predator Prey (PP) | speed-up: 1.7× | Helped |
| Use an FSM implementation instead of standard NUTS | NUTS · Google Stock (GS) | speed-up: 3.5× | Helped |
| Use an FSM implementation instead of standard TESS | TESS · Google Stock (GS) | speed-up: 2.2× | Helped |
| Use an FSM implementation instead of standard NUTS | NUTS · German Credit (GC) | speed-up: 1.2× | Helped |
| Use an FSM implementation instead of standard TESS | TESS · German Credit (GC) | speed-up: 1.0× | No clear effect |
| Use an FSM implementation instead of standard NUTS | NUTS · Biochemical Oxygen Demand (BOD) | speed-up: 0.8× | Hurt |
| Use an FSM implementation instead of standard TESS | TESS · Biochemical Oxygen Demand (BOD) | speed-up: 1.1× | Helped |
| Amortize log-pdf calls in the elliptical slice sampler (ablation) | FSM elliptical slice sampler · Real Estate UCI | speed-up: 2× | Helped |
| Do not amortize log-pdf calls (ablation) | FSM elliptical slice sampler · Real Estate UCI | relative performance: factor 2 | Hurt |
| Bundle the delayed-rejection FSM into a single step instead of four states (ablation) | Symmetric DRMH FSM · Univariate Gaussian | efficiency gain: nearly a 3x efficiency gain | Helped |
| Use the condensed single-state FSM instead of the four-state FSM (ablation) | Symmetric DRMH FSM · Univariate Gaussian | speed-up: near an order of magnitude | Helped |
Surprising findings
- Chronological block ordering was found to be a surprisingly effective heuristic for step bundling.
- On some tasks (e.g., BOD for NUTS, GC for NUTS/TESS) FSM gives no speed-up or slight slowdown (0.8x, 1.0x).
Limitations
- Speed-up depends on the skewness of the per-chain iteration distribution; gains shrink when iteration counts are homogeneous or log-pdf costs dominate differently.
- Theoretical analysis uses a simplified cost model with a single while loop; the general FSM framework is described but only tested on a few algorithms.
- Each state executes all branches under vmap, so the per-step cost grows with the number of states.
Experimental setup
Datasets: UCI Real Estate Valuation Predator Prey Google Stock German Credit Biochemical Oxygen Demand
Benchmarks: Real Estate GP hyperparameter posterior univariate Gaussian DRMH 100-dim correlated mixture of Gaussians NUTS four challenging-geometry distributions for TESS/NUTS
Metrics: wall-clock time effective sample size per second ESS/sec average iterations per sample speed-up ratio
Baselines: BlackJAX elliptical slice sampler BlackJAX HMC-NUTS standard symmetric delayed rejection MEADS CHEES NUTS (non-FSM)
Builds on: elliptical slice sampling No-U-Turn Sampler delayed rejection Metropolis-Hastings JAX vmap BlackJAX
Closest papers by meaning
- Parallelizing MCMC Across the Sequence Length (NeurIPS 2025 · similarity 0.85)
- AutoStep: Locally adaptive involutive MCMC (ICML 2025 · similarity 0.84)
- Involutive MCMC: a Unifying Framework (ICML 2020 · similarity 0.84)
- Accelerating Monte Carlo Tree Search with Probability Tree State Abstraction (NeurIPS 2023 · similarity 0.83)
- Path Auxiliary Proposal for MCMC in Discrete Space (ICLR 2022 · similarity 0.83)
- Accelerating Look-ahead in Bayesian Optimization: Multilevel Monte Carlo is All you Need (ICML 2024 · similarity 0.82)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.