Efficiently Vectorized MCMC on Modern Accelerators
Authors: Hugh Dance
Pierre Glaser
Peter Orbanz
Ryan P Adams
Venue: ICML 2025 spotlightposter (2025)
Keywords: Markov Chain Monte Carlo
MCMC
Parallel Computing
Primary area: probabilistic_methods->monte_carlo_and_sampling_methods

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 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 vmap, 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.

Efficiently Vectorized MCMC on Modern Accelerators
Hugh Dance' Pierre Glaser'
Peter Orbanz?
Ryan Adams?
Abstract
MCMC chains in parallel: one can simply write single-chain MCMC code, and call vmap to turn it into vectorized, multi- chain code that can run in parallel on the same processor. Many state-of-the-ar MCMC libraries have consequently adopted machine learning frameworks with automatic vec- torization tools as their backend.
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 gooblem-lossee 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 simpli- fied model, we derive 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.
speaking.
One limitation with automatic vectorization tools is how they handle control flow. Since all instructions must be executed in lock-step, if the algorithm has while loop, all chains must wait until the last chain has finished its it- erations. This can lead to large inefficiencies for MCMC algorithms that generate each sample using variable-length while loops. Roughly speaking, if vectorization executes, say, 100 chains in parallel, all but one finish after at most 10 steps, and the remaining chain runs for 1000 steps, then about 99% of the GPU capacity assigned to vmap is wasted (and our simulations show the effect can indeed be this dras- tic). For the No-U-Turn Sampler (HMC-NUTS) (Hoffman et al., 2014), this problem is well-documented (BlackJax, 2019; Sountsov et al., 2024; Radul et al., 2020). also af- fects various other algorithms, such as variants of slice sam- pling (Neal, 2003; Murray al., 2010; Cabezas and Nemeth, 2023), delayed rejection methods (Mira et al., 2001; Modi et al., 2024) and unbiased Gibbs sampling (Qiu et al., 2019).
1. Introduction
In this work, we show how to transform MCMC algorithms into equivalent samplers that avoid these synchronization barriers when using vmap-style vectorization. In particular,
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 single algorithm concurrently. They are routinely used to train neural networks (Flax, 2023) and in other scientific applications, e.g., Schoenholz and Cubuk (2021); I); Oktay al. (2023); Pfau al. (2020).
1. We develop a novel approach to transform MCMC algorithms into finite state machines (FSMs), that can avoid synchronization barriers with vmap.
2. We analyze the time complexity of our FSMs against standard MCMC implementations and derive theoret- ical bound on the speed-up under simplified model.
This paper focuses on the use of automatic vectorization for Markov chain Monte Carlo (MCMC) algorithms. Tools like JAX's vmap provide convenient way to run multiple
3. We use our analysis to develop principled recommen- dations for optimal FSM design, which enable us to nearly obtain the theoretical bound in speed-ups for certain MCMC algorithms.
Gatsby Unit, University College London, UK Department of Computer Science, Princeton University, USA. Correspondence to: Hugh Dance <hugh.dance. 15@uul..c..k>>.
4. We implement popular MCMC algorithms as FSMs, including Elliptical Slice Sampling, HMC-NUTS, and MH - demonstrating speed-up ups of up to an order of magnitude in experiments.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Efficiently Vectorized MCMC on Modern Accelerators
2. Background and Problem Setup
lation), vmap has been adopted by major MCMC libraries such as NumPyro and BlackJAX.
In this section, we briefly review how MCMC algorithms are vectorized, explain the synchronization problems that can arise, and formalize the problem mathematically.
2.2. Synchronization Problems with While Loops
Control flow (i.e., if/else, while, for, etc.) poses a challenge for vectorization, because different batch mem- bers may require different sequence of instructions. vmap solves this by executing all instructions for all batch mem- bers, and masking out irrelevant computations. As a consequence, if sample contains a while loop, then vmap ( (sample) will execute the body of this loop for all chains until all termination conditions are met. Until then, further instruction can be executed. As result, if there is high variation in the number of loop iterations across chains, running Algorithm with vmap (sample) introduces a synchronization barrier across all chains: at every iteration, each chain has to wait for the slowest sample e call.
2.1. MCMC Algorithms and Vectorization
MCMC methods aim to draw samples from a target distri- bution и (typically on a subset of IRd) which is challenging to sample from directly. To do so, they generate samples (#1,...,<E. } E Rdx= from a Markov chain with transition kernel P and invariant distribution x, by starting from an ini- tial state my E IRd and iteratively sampling #++1 ~ P(-lx;)- For aperiodic and irreducible Markov chains, as n - 0o the samples will converge in distribution to æ (Brooks et al., 2011). In practice. P implemented by a deterministic function smmll, which takes in the current state æ; and a pseudo-random state r, E N, and returns new states x,+1 and ri-1- This procedure is given in Algorithm 1.
This issue arises in practice, because a number of impor- tant MCMC algorithms have while loops in their sample implementations: such variants of slice sampling (Neal, 2003; Murray et al., 2010; Cabezas and Nemeth, 2023), delayed rejection methods (Mira et al., 2001; Modi et al., 2024), the No-U-Turn sampler (Hoffman et al., 2014), and unbiased Gibbs sampling (Qiu et al., 2019).
Algorithm 1 MCMC algorithm with sample function
1: Inputs: sample Eo. seed ro
2: for i {l,..., n do
3: generate F; Ti f sample 1;-1)
4: end for
5: return F1....,,r, Zn
Formalizing The Problem. Here we formalize the problem through a series of short derivations. These will be made precise in Section 4. Suppose of Algorithm 1 using a sample e function that has while loop. Let Ne,, denote the number of iterations required by the jith chain to obtain its ith sample. If the while loop has a variable length, NN. is a random number. Due to the synchronization problem described above, the time taken to run vmap ( (sample) at iteration is approximately propor- tional to the largest No, out of the chains, maxjsm Neg- If we ignore overheads, the total runtime after samples, Co(n), is then roughly proportional to
Practitioners commonly run Algorithm on m different initializations, producing m chains of samples. Given an implementation of sample: e: T2d N + Rdd N, one way to do this is to use an automatic vectorization tool like JAX's vmap'. vmap takes sample program, vmap (sample): Rdxm x N + ddxm x N which operates on a batch of inputs collected into tensors
vmap 'dd version
input and returns new
3;-1 := (Z:-1,1, ... ((1--1))...3311mm) E dd Ā--1 == (fi-1.1,..... (.--1.,......11... 1mm) Nm
and returns the corresponding outputs from sample. One can therefore turn Algorithm 1 into multi-chain algo- rithm by simply replacing sample with vmap (sample) and replacing (z) by (io,Fu) Fu) (see Algorithm 6 in Ap- pendix B). Under the hood, vmap transforms every instruc- tion in sample (e.g., dot product) into a corresponding instruction operating on a batch of inputs (e.g., matrix-vector multiplication); that is, it 'vectorizes' sample. These in- structions are executed in lock-step across all chains. Using vmap usually yields code that performs as well as manually- batched code. For this reason, well as its simplicity and composability with other transformations like grad (for compi-
Co(n) X
Σ max Nog
35m
By contrast, if the chains could be run without any synchro- nization barriers, the time taken would instead be
(.()
X max
N
S
The key difference is that the maximum is now outside the sum. This reflects the fact that when running independently, we only have to wait at the end for the slowest chain to collect its n samples, rather than waiting at every iteration. Clearly (m) & (n)). Significant speed-ups are obtain- able by de-synchronizing the chains when C. (n) < Co(n). If each Ni,, converges in distribution to some P. as we
automatic differentiation) and jit (for
We use JAX and its vectorization map vmap throughout, since this framework is widely adopted. Similar constructs exist in TensorFlow (vectori i zed map) and PyTorch (vmap).
2

---

Efficiently Vectorized MCMC on Modern Accelerators
Siice shrinks per sampre
Aug shrinies per thein n semples
per chain is contained in the interval [5,8], and after 10,000 samples this shrinks to [6.2, This implies that if we could vectorize the algorithm without incurring synchro- nization barriers, we could improve the Effective Sample Size per Second (ESS/Sec) by up to 3-fold. We will see that for other algorithms (e.g., HMC-NUTS and delayed rejection) the potential speed-ups are much larger than this.
f-100
1.0
Avg of Dax 0004 cha not
Ang Mes (320 thaina
10
n=10000
n-1000
Daiiiiy
000
0.2
6.0 siice shrinis
slice shrimks
3. Finite State Machines for MCMC
Figure 1. Statistics on the elliptical slice sampler for Gaussian Pro- cess Regression on the Real Estate Dataset (Ych, 2018). LHS: his- togram of the number of slice shrinks per sample. RHS: smoothed histogram of the average number slice shrinks per sample across m 1024 chains, after E {100, 1000, 10000) samples.
In this section we present an approach to implementing MCMC algorithms, which avoids the above synchronization problems when vectorizing with vmap. The basic idea is to break down the transition kernel (samp e) into a series of smaller 'steps' which avoid iterative control flow like while loops and have minimal variance in execution time. We then define runtime procedure that allows chains to progress through their own step sequences in de-synchronized fash- ion. To do this in principled manner, we adapt the frame- work of finite state machines (FSMs) (Hoperoft et al., 2001). An FSM is a 5-tuple (S, 2, 5, B, F), where is a finite set of states, Z is finite input set, 5 : Sx + Ss a transition function, B E S is an initial state, and F is the of final states. We deviate from the usual definition by allowing each state to be function S : 2 + Z that updates inputs dynamically². Below we show how to represent different sample e functions as an FSM
draw more samples, and an appropriate central limit theo- rem holds (so that the sums in (1) and (2) behave like scaled means), we can expect for large enough that:
(m) = Op E max Noo.s) 35m
(3)
(n) Op(max n ENo...) = Opp( jmm
(4)
Hid
where the last equality assumes (No,, ..., No De-synchronizing the chains will result in large speedups if
~ FN.
E max No..j J.m
>>>
EN
(5)
3.1. sample-to-FSM conversion
We will see that Equation (5) holds in various situations.
At high-level, we construct the FSM of an MCMC algo- rithm as follows. The states S = (S,, are chosen as functions Z + Z which execute contiguous code blocks of the algorithm. The boundaries of each code block are delineated by the start and end of any while loops. For example, Si executes all lines of code before the first while loop, S2 executes all lines of code from the beginning of the first while loop body to either the start of the next while loop or the end of the current while loop, and so on. The inputs 2 E 2 are all variables used in each code block. (e.g. current sample æ, proposal x', log-likelihood log p(x))) whilst the transition function 8 selects the next code block to run according the relevant loop condition and the re- ceived output 2' after executing the current block (e.g. if a while loop starts after the block S₁ executes, 0 will check this loop's 's condition using the output 2' = =S(()) Below we make this construction procedure more precise for three kinds of sample functions which cover all MCMC algo- rithms considered in the present work: (i) functions with a single while loop, (ii) functions with two sequential while
Example. Consider the elliptical slice sampling algo- rithm (Murray et al., 2010) which samples from distribu- tions which admit a density with Gaussian components, p(x) (x)) 00, E). Its transition kernel (see Algo- rithm& 8 in Appendix B) draws each sample by (i) generating a random ellipse of permitted moves and an initial proposal, and (ii) iteratively shrinking the set of permitted moves and resampling the proposal from this set until it exceeds a log- likelihood threshold. The second stage uses while loop which requires a random number of iterations.
On Figure I we display results when implementing this algorithm in JAX to sample from the hyperparameter pos- terior of a Gaussian process implemented on a regression task using a real dataset from the UCI repository (details are in Section 7). On the LHS we can see the distribution of the number of while loop iterations (i.e. slice shrinks) needed to generate sample. While the average is ~6, the average of the maximum across 1024 chains is >> which implies E max Nn.j R 3EN.,,1 On the RHS, we can see that the differences across the chains do balance out as more samples are drawn, which suggests a certain central limit theorem may hold. In particular, after just 100 samples the distribution of the average number of iterations
This essentially results in a so-called 'Finite State Transducer' with a different runtime than typical.
"Typically x Rdd so 2 is finite. However, in practice all samples lie a finite subset I2d (e.g. 32-bit floating points).
3

---

Efficiently Vectorized MCMC on Modern Accelerators
loops, and (iii) functions with nested while loops. An automated construction process for more general programs using a toy language is given in Appendix B.
loops, break down sample into B1, B2, B3, where B1 (resp. B3) is the code before (resp. after) the outer while loop and B2 is the outer while loop body. As B2 is a single- while loop program, it admits its own FSM F.. Building the final FSM of samp then informally consists in obtaining a first, "coarse" FSM by treating B3 as opaque, and then refining by replacing B2 with its own FSM. The full con- struction process is given in Appendix B. The resulting FSM is provided for the case of NUTS (Figure 3, bottom-right panel), which-in its iterative form the outer while loop to determine whether to keep expanding Hamiltonian trajectory, and the inner while loop to determine whether to keep integrating along the current trajectory.
code block
S1
S1
while... do {
code block 2
=: S2
S2
end while
code block
S3
S3
3.2. Defining the FSM Runtime
Going forward, for convenience we assume the transition function 5 takes in and returns the label of each block Sk, rather than Să itself. Now, given a constructed FSM, in Algorithm 2 we define function step, which when executed performs single transition along edge in the FSM graph. For reasons made clear shortly, augment step to return flag indicating when the final block is run.
Figure 2. FSM of an MCMC algorithm with a single while loop.
The Single While Loop Case. The simplest case is a sample method with a single while loop. This covers (for example) elliptical slice sampling (Murray et al., 2010) and symmetric delayed rejection et al., 2001). In this case, we break sample e down into three code blocks B1, B2, B3 (one before the while loop, one for the body of the while loop, and one after the while loop) and the termination condition of the loop. This is shown on the LHS of Figure 2. Using these blocks, we define the FSM as (8, 2, 6, F). where (1) Z is the set of possible values for the local variables of sample (e.g., the current sample æ, seed and proposal etc.), (2) S (Si, S2, 83) } is set of three functions Z + Z. where for each A E 21, 3). (²)z) runs By on local variables = and returns their updated value, (3) for each e 13 and z E Z, the transition function 8(k,, 2) checks the while loop termination condition using 2, and returns S2 if False and S3 if True, (4) B S and, finally, (5) F Sz. The RHS of Figure illustrates the resulting FSM diagram. Note there is an edge Sy + S between states if and only if 5(Sx, 2) = S for some 2.
Setropolis-Hastinge (Mira
Algorithm 2 step function for FSM
Inputs: algorithm state k. variables z
I: set sSample +- 1(k: K}
2: z + switch(k, [[run S (z)),.... (..))............. frun Sx
3: + switch(k, [frun 6(1, frun 5 (K, 2)D
4: return (k, z, aammll)
Here swi witth uses the value of k to determine which branch to run. If we start from some input (k, (0, 20) and call step repeatedly, we transition through sequence of blocks of sample, until we eventually reach the terminal state. At this point, a sample is obtained (as indicated by i isSample-True) We can use this function to draw n samples. by (i) adding a transition from the terminal state Jf Sx back to the initial state B So, and (ii) defin- ing wrapper function which iteratively calls step until i ssampple=True is obtained n-times (see Algorithm 3).
Two Sequential While Loops. We break down sample l into two blocks: B1 contains all the code up to the second while loop. Bg contains the remaining code. In this case, B and B2 are now single while loop programs, and thus can both be represented by FSMs F, and F2 using the above rule. The FSM representation of sample can then be obtained by "stitching together" " F1 and F2. The full construction process is in Appendix B. The resulting FSM provided for the case of the Slice Sampler (Figure 3, top-right panel),
Algorithm 3 FSM MCMC algorithm
1: input: initial value æo # samples
initialize: Z init(x)), list() B list()
3: Set t 0 and Is
4: while /. < n do
5: (k: 2, isSample) + step(k,z)
which contains two* while loops: one for expanding the
6: append current sample value æ stored in = to X
slice, and one for contracting the slice.
7: append isSample e to B
Two Nested While Loops. In the case of two nested while
update sample counter t <- t + isSample
end while
'For ID problems, the slice expansion loop can he broken into two loops (for the upper and lower bound the interval).
10: return X[B]
4

---

Efficienlly Vectorized MCMC on Modern Accelerators
Delayed Rejection MH
Slice Sampling
INIT
PROPOSE
DONE
INIT F
EXPAND
INIT S
SHRINK
DONE
S1
2
S3
Si
S₂
S3
S4
S5
INIT
SHRINK
DONE
INIT
DOUBLE
INTEGRATE
CHECK
DONE
S
2
S3
S
S2
S3
S4
S5
Elliptical Slice Sampling
HMC-NUTS
Figure 3. Finite state machines for the sample function of different MCMC algorithms: The symmetric delaycd-rejection Metropolis- Hastings algorithm (Mira et al., 2001), elliptical slice sampling (Murryy et al.. 2010), (vanilla) slice sampling (with single slice expansion loop) (Neal, 2003), the No-U-Ture sampler for Hamiltonian Monte Carlo (Hoffman et al., 2014).
Both Algorithm 3 and Algorithm 1 call sample n-times and can be easily vectorized with vmap. For Algorithm 3, we just call vmap on step and modify the outer loop to terminate when all chains have collected samples each (see Algorithm 7 in Appendix B). The crucial difference is: (i) by defining 'step' to be agnostic to which block each chain is currently on, Algorithm enables chains simulta- neously progress their independent block sequences, and (ii) by moving all while loops to the outer layer, Algorithm 3 only requires chains to synchronize after samples.
for chain j. We assume the cost of executing each block Sp with vmap is cr(m) for chains (here the dependence on m reflects underlying GPU vectorization efficiency). The average cost per sample of Algorithm after n samples, when vectorized for m chains, is then
Σ max Ni. jElm]
Co(m, n) = Ao(m) + Bo(m)
(6)
where Ao(m) == ej (m). Bo(m) ca (m), and Nij # calls of the loop body S₃ for chain j and sample The max reflects the synchronization barrier across chains.
4. Time Complexity Analysis of FSM-MCMC
In this case, for the FSM (vmap'ee ed Algorithm 3), the cost of calling each block S is Σ cy(m), since vmap (step) executes all switch branches for all chains. We will assume this cost can scaled by some a [maake]]]] ex (mm) E C* (ma), 1]. since we will later re- design step to reduce a < This gives an average cost per sample for chains after samples of,
One limitation with our FSM design is the step function relies on switch determine which block to run. When vectorized with vmap branches in swi swithhhhc tch are evaluated for all chains, with irrel- evant results discarded. This means the cost of single call of step is the cost of running all state functions Si,.., SK- To obtain a speedup, the FSM must therefore sufficiently decrease the expected number of steps to obtain n samples from each chain. this section we quantitatively derive con- ditions under which this occurs in the simplified setting of an MCMC algorithm with single while loop. This enables us to subsequently optimize the FSM design in Section 5.
equivalent transformation, all
Cp(m, n) Ap(m) + Bp(m) max Σ N
(7)
jE(m] n
i=1
where Ap(m) a(K - ()(c(((m)) + (m)), Bp(m)
a(c; (m) c-xx cmx(m))) and (m) == = Ej (m().
Let Xt,, be the 1th random sample obtained by chain j. If the joint sequence (X,,j= Ni,)221 an appropriate Markov chain, we have the following concentration result for Cu, CF. which formalizes derivations in Section 2.2.
4.1. Long Run Costs of FSM and Non-FSM MCMC
To this end, consider sample e function with single while loop. Suppose it is broken down into K blocks (S1,...kx) }. by our FSM construction procedure, where Sş (for some E [K] executes the body of the loop. Note our procedure yields K < 3 for one while loop, but we relax this here to analyze the effect of the number of blocks on performance. single call of vmap (sample) executes St once if + k. and maxy N, times if I = k, where N₁ = # loop iterations
Theorem 4.1 (Long Run Costs). Let Ni.; E [0, B), E X' IRd, (Xi,NNiJ)a be a Markov Chain with station- ary distribution π with absolute spectral gap 1- E [0, 1). Then with probability 1 we have the inequalities:
[Co(m, n)) - Co(m)) < MB In(2/6)
(8)
CF(m,,) Cr(m)) < MBF(m)nn In(2m/6) (9)
5

---

Efficienlly Vectorized MCMC on Modern Accelerators
sid Where, for some (N1..., Nm) ~ NN the long run costs are Co(m) == (m) + (m max N jE[m]
Distributions Stewnoss
Distributions Siowncss 10
100
(10)
Cp(m)= == AF (m) + B (m)))
(11)
Bern 275) Poics(3)
BernjC, C3)
Polsc/0.011
Betato
Beeatt.
Gamme(4, 1)
GammarO 04, 11
The result says that as n - do, the cost per sample of the standard MCMC design converges to the expected time for the slowest chain draw sample, whilst the FSM design converges to the expected time for single chain draw sample, scaled by the additional cost of control-flow in step (i.e. Ag (m) and (m)))). Both rates are On )).
0-
6
+00
B00
1000
200
400
6úc
800
1000
cheins (m)
chains (m)
Figure 4. R(m) max ge(mn) N2 for different distributions with
Edd
skewness 771 & 1 (LHS) and m ss 10 (RHS).
5.1. Step Bundling to Reduce K
4.2. The Relative Long- Run Cost of the FSM
Given an FSM with step defined by code blocks Si,...,Kk and transition function é, one can 'bundle' multiple steps together using a modified step function, bundled_step which replaces the switch over the al- gorithm state in step with a series of separate condi- tional statements. This is shown in Algorithm 4 for an example with two state functions. Note that under the "run all branches and mask" behaviour of vmap, vmap (step) and vmap (bundl ed ed_step) have the same cost. How- ever, whenever bundled_step runs S₁ and â returns k = 2, it immediately also runs Sg. This reduces the "effec- tive' number of states K and/or the overall number of steps needed to recover a sample, increasing efficiency. In princi- ple, the block ordering can be optimized for sequences that are expected to occur with higher probability. However, we found chronological order a surprisingly effective heuristic.
Using the form of the constants in Theorem 4.1, the ratio of the long-run costs, (mm) == Co (m), is given by
(m)) + (m) R max]e[m] (ma N₂ E(m) a(c_s(m) + (3(m)) 1+ EN1)
(12)
In Proposition A.3 in Appendix A we prove that:
E [aa N, ENN
E(m)= < R(mm) ==
(13)
and that this bound is tight. We refer to R(m) as the 'theoret- ical efficiency bound' for the FSM. Note from Equation (12) that minimizing a and K improves the efficiency E(m) of the FSM. In Section 5 we introduce two techniques to mini- mize and K in practice, which enable us to nearly obtain the efficiency bound R(m) for certain MCMC algorithms.
A Algorithm 4 bundl ed_step for FSM with S2
Scale of Potential Speed-ups. The size of R(m) depends (only) on the underlying distribution of N₁ (since N; =a N, V [m]). Whenever N1 is sub-exponential, it is known that R(m)= = ((n(()) (Vershynin, 2018). Although this implies a slow rate of increase in m. R(m) can be still be very large for small values of m. For example, if N5/B- ~ Bern(p) (i.e., one either needs zero or B iter- ations to get a sample), then R(m) = (1 - (1 p/p and converges to 1/p exponentially fast as m increases. For small p this can be very large even for small m. In gen- eral R(m) is sensitive to the skewness of the distribution: distributions on [0, B] with zero skewness have R(m) = 2, whilst distributions with skew of 10 can have R(m)= = 100: see Figure+ 4. Intuitively, these are the distributions where chains are slow only occasionally, but at least one chain slow often. In such cases, our FSM-design can lead to enormous efficiency gains, as we show in experiments.
Inputs: algorithm state kt variables =
I: set issample + 1,k 2}
2: if k: = then
3: run block S with local variables z
4: update state <- 5(1,2)
end if
6: if 1: 2 then
7: run block S2 with local variables
8: update state k + 8(2,z)
9: end if
10: return 2, (k.,,isSample)
5.2. Cost Amortization to Reduce a
If a function g is called on a variable E 2 inside M < K state functions, a single call of vmap (step) (or vmap (bundled _dd_stp)))) will execute g M-times. To prevent this, we (i) augment step to return a flag doComputation indicating if g executed in the next code block, and (ii) define a new step function amort ioortizee aeed__seeeeeeeeeeeeeeeeeeeeeeeee around step which calls st once, and executes g if doComput ion=TTrue The resulting step function is shown in Algorithm (note if g is called in step Si, we assume it already pre-computed in 2).
5. Optimal FSM design for MCMC algorithms
In this section we provide two strategies to modify the func- tion step to (effectively) reduce a and K. These strategies enable us to develop MCMC implementations which nearly obtain the theoretical bound R(m) in some experiments.
6

---

Efficiently Vectorized MCMC on Modern Accelerators
Note the function g is now only called once per step when using vmap on amorti zed_ step. In partieular, if the executions of g cost (in total) 3 E [0, 1J fraction of the to- tal cost of all blocks, amortization will reduce this fraction to M + (1 (3) We indeed observe this when amortizing g = logp for the elliptical slice sampler in Section 7.1. In that scenario, M- = and 3 1 as the log-likelihood is very expensive, resulting in x speed-up over no amorti- zation. One issue is that amortization only allows steps to be bundled which do not require calls of g. We adapt the algorithm to handle this in Appendix B.
single batched primitive that cannot be broken down by their algorithm. Additionally, our method uses the fewest (while-loop-free) blocks possible and allows for control flow within blocks (which we showed is crucial for optimal per- formance under the "run and mask" paradigm). By contrast, their algorithm does not allow blocks to contain any control- flow, yielding many small blocks. We also obtain provable speed-ups for MCMC via theoretical analysis.
7. Experiments
The following experiments evaluate FSMs for different MCMC algorithms with while loops against their stan- dard (non-FSM) implementations, as well as other MCMC methods in one experiment. All methods consist of sin- gle chain MCMC algorithms written in JAX, turned into multi-chain methods with vmap, and compiled using jit. All experiments are run in JAX on an NVIDIA A 100 GPU with 32GB CPU memory. Code can be found at https: / //github. com/hwdance/ jax- fam-s-m-eee mcmc.
Algorithm 5 amort i zze_ssepp for FSM with function g
I: Input: Algorithm state k. variables 2
2: (R:, 2, z,sismple e, doComputation)-- 2k)
3: if doComputat on then
4: Unpack state (='
5: Do computation 0 4 (2)
6: Re-pack state= 2 + (2)
7: end if
8: Return 2k, issampee)
7.1. Delayed-Rejectinn MH on a Simple Gaussian
6. Related work
We first illustrate basic properties of the FSM conversion on a toy example. The MCMC algorithm used is symmet- ric Delayed-Rejectionn Metropolis Hastings (DRMH) (Mira et al., 2001), which is a simple example of a delayed rejec- tion method. Symmetric DRMH modifies the Random-Wan Metropolis- Hastings algorithm by iteratively re-centering the proposal distribution on the rejected sample and resam- pling until either acceptance occurs or maximum number of tries M is reached. To ensure detailed balance, the accep- tance probability is adjusted to account for past rejections.
FSMs in Machine Learning. Previous work in machine learning has used the framework of FSMs to design image- based neural networks (Ardakani et al., 2020) and Bayesian nonparametric time series models (Ruiz et al., 2018), well as extract representations from Recurrent-
(RNNs) (Muškardin et al., 2022: Cechin et al., 2003; Tiño et al., 1998; Zeng et al., 1993; Koul et al., 2018; Svete and Cotterell, 2023). To our knowledge, our work is the first that uses FSMs as a framework to represent MCMC algorithms, and design novel implementations.
Experimental setup. We implement symmetric DRMH using (vmap' ed) Algorithm I (baseline) and Algorithm 3 (ours) to sample from univariate Gaussian (O 1), vary- ing the number of chains. We use a (x, 0.1) proposal dis- tribution with M = 100 tries per sample and draw 10,000 samples per chain. Although DRMH has 3 state functions by default (see Figure 3), the INIT and DONE states are essentially empty, and so just a single state function can be used for the while loop body. This means an appropriate FSM implementation should be able get close to R(m), up to overheads. To illustrate the importance of designing FSMs with as few (effective) state functions as possible, we implement an FSM which unrolls the while loop body into four different state functions, as well as a (condensed) FSM with a single state function (i.e. effectively uses bundling).
Efficient MCMC on Modern Hardware. Given the recur- sive nature of HMC-NUTS, previous work has reformulated the algorithm for compatibility with machine learning frame- works that cannot naively support recursion (Abadi et al., 2016; Phan et al., 2019; Lao et al., 2020). However, these implementations do not address synchronization inefficien- cies caused by automatic vectorization tools. Our work bears some similarities with general-purpose algorithm proposed in the geigh-Perfomance-Cooe Computtng literature for executing batched recursive programs (Radul et al., 2020). Both their method and ours breaks programs down into smaller blocks for efficient vectorization, but there are sev- eral major differences. Our approach provides a recipe for implementing single-chain iterative MCMC algorithms and is fully compatible with automatic vectorization tools like vmap. In contrast, their algorithm is designed for recursive programs and requires code which is already batched. Cru- cially, to avoid synchronization barriers due to while loops, this code cannot have been batched with a vmap-equivalent vectorization tool, since vmap converts while loops into
Results. As the number of chains m increases, the FSM implementations increasingly outperform the standard im- plementation (see Figure 6). This reflects the increasing synchronization cost of waiting for the slowest chain. The speedups are near an order of magnitude when m = 1024
7

---

Efficiently Vectorized MCMC on Modern Accelerators
Avg iterations per sample
Walltime
ESS per second
FSM relative efficiercy gain
2.0k
2.0K
F3M
RlackjXX sice
21
terations -5
1.5k
ESSibecond
1.5k
Efficioncy
20
1.0%
1.0k
20
torsisampl
0.5k
0.5k
2.1
walkome
waltiiee lodt in 24
0.0k
2-2
22
22
2/ chalns
29 20
24
230
20
24
24
24
2:
20240
20
2 24 2' 22
2P.12
24 22 2b 2/ chains
20 2i0
# chaans
chains
Figure 5. Average results using 10 random seeds (standard deviations too small to show) from drawing 10k posterior samples for the covariance hyperparameters (T, 0, .) of a Gaussian Process (a)) = f(x + on the Real Estate UCT dataset (n 411, d 6). Blue BlackIAX elliptical slice, Red = FSM clliptical slice sampler (red). LHS: the average number of sub-iterations (i.c., ellipse contractions) needed to draw a single sample increases from (1 chain) to 18 (1024 chains) for the standard implementation due to synchronization barriers, but remains constant for our FSM. Middle two plots: The FSM can run ~3x faster by avoiding synchronization barriers, as shown by the Walltime (left-middle) and ESS/S (right-middle). RHS: the ratio of average iterations per sample (i.e. A(m)) (green) bounds the obtainable 'efficiency gain' using our FSM, but is nearly obtained in relative walltime when amortizing log-pdf calls (red).
Walltime
Effective Samples per Second
and assume (V(0, k(2,z") 72 eexp(->21 - 2'1))). We use Normal priors O, ~ (0, 1), (so the ellipse is drawn using N T)) and use the sampler to draw 10k samples per chain from the posterior p(a, T, A(D.), for varying # chains m.
GP(O,)) k) with kernel
22
Standard DR
2
FSM UK
FSM DR (condonsed)
Spatin plan
22
22
27
23
2-
2-5
Results are shown in Figure 5. As expected, the BlackJAX implementation suffers from synchronization barriers at ev- ery iteration due to using vmap with while loops: its average number of iterations per sample increases roughly logarith- mically from 6 (1 chain) to 18 (1024 chains), whereas the FSM implementation remains constant. As a result, the FSM significantly improves walltime and ESS/second per- formance (Figure 5/middle). For instance, when 1024 chains are used, the FSM reduces the time to draw 10k samples per chain from over half hour to about 10 minutes. The effi- ciency gain can be measured by the ratio (BlackAAX/FSM) of wall-times (shown in Figure S/right). As expected. FSM efficiency increases with the number of chains. The greatest efficiency gain occurs precisely where the best ESS/second can be obtained via GPU parallelism. The analysis in Sec- tion 4 shows that the efficiency gain is upper-bounded by the ratio of average # iterations per sample for both methods (i.e., R(m)). This bound is almost achieved here by amor- tizing log-pdff calls. This is because (i) roughly 80% of the time is spent in the iterative 'SHRINK' state, and (ii) log-pdf calls dominate computational cost here, so amortizing them ensures the state function cost is similar to the standard im- plementation. Since the log-pdf is needed in two states, not amortizing it results in two log-pdf calls per step, and so we lose roughly factor 2 in relative performance (orange line in Figure 5). These results change with data set size, which determines the cost of log-pdf calls (see Appendix B).
2' 2
2
2' 2
21
2' 2' 2 chalnss
chains
Figure 6. Mean and standard deviation walltimes (LHS) and ESS (RHS) of the symmetric
Rejoyccinnn Algorithm (Mira et al., 2001) using Algorithm and our FSM implementation Algo- rithm 3 (both with vmap), on a univariate Gaussian (10 random seeds). FSM uses stee with the while loop body (PROPOSE in Figure 3) split into 4 states, whilst FSM-condensed uses single step. Using a single step leads to a 3x efficiency gain.
for the condensed FSM. Bundling sees nearly a 3x efficiency gain and enables no performance loss when m = 1 and there no synchronization barrier. Note Standard DR tracks the profile of the estimated E|max,e]mm] N. whilst the FSMs track the profile of [[ (which is flat). This implies the condensed FSM (i.e. with step bundling) has been able to approximately obtain R(m) up to constant factor.
7.2. Elliptical Slice Sampling on Real Estate Data
The elliptical slice sampler (introduced in the example in Section 2.2) has single while loop, resulting in three state functions (see Figure 3). We compare BlackJAX's X imple- mentation to our FSM implementation.
Experimental setup. We apply the sampler to infer poste- riors on covariance hyperparameters in Gaussian Process Regression, using the UCI repository Real Estate Valuation dataset (Yeh, 2018). This dataset is comprised of n = 414 input and output pairs Dn, = {æ,, 3 Wi where is the house price of area i, and x; E R are house price predic- tors including house age, spatial co-ordinates, and number of nearby convenience stores. We model y = f(x) +
7.3. HMC-NUTS on High-Dimensional Correlated MoN
The NUTS variant of Hamiltonian Monte Carlo (Hoffman et al., 2014) adaptively chooses how many steps of Hamil- tonian dynamics to simulate when drawing a sample, by
8

---

Efficiently Vectorized MCMC on Modern Accelerators
Contours of Log-PDF
Integration Steps per Sample (1 chain)
Integration Steps per Sample (500 chains)
Effective Samples per Minute
HMC-NUTS MALA
1.0
1.0
Standard -UTT r'M HIMC-NUTS
720
10
0.8
0.8
Proboblity
0.6
Probbbiiit
0.6
5500
0.4
0.4
2G3
-10
0.2
0.2-
96
133 6
14:00
47
-20
0.0
0.0
-00
-10
0
10 20
0
200
400
600
900
1000
200
400
600
800
1000
5
20
200
500
X1
#stops
#stops
#cchans
Figure 7. Left: contours of the first two dimensions of correlated mixture of Gaussians, along with a single chain of HMC-NUTS and MALA. Middle Left: Histogram of the number of integration steps taken per sample for single NUTS chain. Middle Right: histogram of the maximum number of integration steps taken per sample across 500 chains. Right: Effective samples per minute for the standard BlackJAX HMC-NUTS implementation, and our FSM implementation of HMC-NUTS (average and standard error bars from sceds). The FSM achieves speed ups of nearly an order of magnitude for 100 chains, and more than half an order of magnitude for 500 chains.
checking whether the trajectory has turned back on itself or has diverged due to numerical error, Its iterative implemen- tation in BlackJAX involves two nested while loops: An outer loop that expands the proposal trajectory, and an inner loop that monitors for U-turns and divergence. Converting these while loops into an FSM using our procedure results in five states (Figure 3). We again compare BlackJAX's implementation to our own (using vmap for both methods).
Table 1. Left: ESS/sec for haselines on Predator Prey (PP), Google Stock (GS), German Credit (GC), Biochemical Oxygen Demand
(BOD) distributions. Right: FSM speed ups for NUTS/TESS.
ESS/sec (non-FSM)
FSM Speed-up
Dist.
MEADS
CHEES
NUTS
TESS
NUTS
TESS
PP
1.5
0.02
2.3
1.5x
1.7x
GS
480.8
60.2
0.15
1116
3.5x
2.2x
GC
141.5
199.0
1.71
58.95
1.2x
1.0x
Experimental setup. We implement NUTS dimensional correlated mixture of Gaussians (p = 0.99), with the mixture modes placed along the principal direction at (-10- 1, 0, 10 1). We use pre-tuned step-size with acceptance rate ~ 0.85 and identity mass matrix. We draw n = 1000 samples per chain and vary # chains m.
100-
BOD
64.85
247.0
0.56
2978
0.8x
1.1x
pling on the transformed distribution TT and then pushes generated samples through T-1 to recover samples from म. TESS achieves particularly good results on distributions with 'funnel' geometries, on which gradient- -based methods like HMC-NUTS tend to struggle (Gorinova et al., 2020). TESS is similar to the elliptical slice sampler and so the FSM has the same 'single loop' structure (see Figure 3).
Results. The trajectory of a single NUTS chain (1000 samples) are displayed on the LHS of Figure 7 (one dot one sample). The typical distance traveled by NUTS is small (few integration steps), with the occasional large jump (many integration steps) when the momentum sample aligns with the principal direction. In particular, the probability a sample requires less than 20 integration steps is ~ 0.95 and needs 1000 steps 0.01, but the probability that at least one chain needs more than 1000 steps is ~ 0.99 (Middle Figure 7). This results in FSM speed-ups of nearly an order of magnitude for m = 100, and about half an or- der of magnitude for mn = 500 (RHS Figure 7). Note that one can avoid sychronization barriers and obtain very high ESS/Sec using a simpler algorithm like MALA, but this fails to explore the distribution (LHS Figure 7).
Experimental setup We examine the speed-ups using FSMs for TESS and NUTS on the four benchmark distributions in Cabezas and Nemeth (2023), which are chosen for their chal- lenging geometries. As baselines we use two otate-of- adaptive HMC variants (MEADS (Hoffman and Sountsov, 2022), and CHEES (Hoffmman et al., 2021)). We expect NUTS to perform poorly on these problems (as observed in Cabezas and Nemeth (2023)), but implementing it lets us at least examine the FSM speed-ups. For all methods, we average results over 128 chains of 1000 samples, with hyperparameters pre-tuned over 400 warm-up up steps.
Results are in Table 1. TESS-FSM improved ESS/sec over TESS in all cases (in 2/4 cases by 2x) and in 3/4 cases improved the best ESS/sec. NUTS performed poorly as expected, but the FSM improved ESS/sec in 3/4 cases, and > 3x in one case. Speed-ups for NUTS and TESS were larger for tasks with an expensive log-pdf (PP, GS). This is because most time is spent in the loop states where the log-pdf is called, and a more costly log-pdf reduces the gap between the cost of executing these states for non- non-FSMs (ck) and FSMs E/ (g) (see Section 4), improving efficiency.
7.4. Transport Elliptical Slice Sampling and NUTS on Distributions with Challenging Local Geometry
Transport elliptical slice sampling (TESS), due Cabezas and Nemeth (2023), is state-of-the-art variant of elliptical slice sampling designed for challenging local geometries. It uses normalizing flow T to 'precondition' the distribution v into an approximate Gaussian, does elliptical slice sam-

---

Efficiently Vectorized MCMC on Modern Accelerators
Acknowledgements
Gorinova, M., Moore, D., and Hoffman, M. (2020). Auto- matic reparameterisation of probabilistic programs. In International Conference on Machine Learning, pages 3648-3657, (6657.PM PMLR.
HD, PG and PO are supported by the Gatsby Charitable Foundation. This work was partially supported by NSF
OAC 2112201.
Hoffman, M., Radul, A., and Sountsov, P. (2021). An adaptive-meme scheme for setting trajectory lengths in hamiltonian monte carlo. In International Conference Artificial Intelligence and Statistics, pages 3907-3915.
Impact Statement
This work improves the execution efficiency of MCMC algorithms. Our method has the potential to significantly reduce time, cost, and energy expenditure for a range of scientific applications.
PMLR.
Hoffman, M. D., Gelman, A., et al. (2014). The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo. J. Mach. Learn. Res., 15(1):1593-166 1623.
References
Abadi, M., Agarwal, A., Barham, P., Brevdo, E., Chen, Z.,, Citro, C., Corrado, G. S., Davis, A., Dean, J., Devin, M., et al. (2016). Tensorflow: Large-scale machine learning on heterogeneous distributed systems. arXiv preprint arXiv: 1rXiv:1603.04467
Hoffman, M. D. and Sountsov, (2022). Tuning-free gener- alized hamiltonian monte carlo. International confer- ence artificial intelligence and statistics, pages 7799- 7813. PMLR.
Hoperoft, J. E., Motwani, R., and Ullman, J. D. (2001). Introduction to automata theory, languages, and computa- tion. Acm Sigact News, 32(1))60--5.
Ardakani, A., Ardakani, A., and Gross, W. (2020). Training linear finite-state machines. Advances in Neural Informa- tion Processing Systems, 33:7173-7183.
Koul, A., Greydanus, S., and Fern, (2018). Learning finite state representations of recurrent policy networks. arXiv preprint arXiv:1811.122.
BlackJax (2019). Sample with multiple chains in parallel.
Brooks, S., Gelman, A., Jones, G., and Meng, X.-L. (2011). Handbook of markov chain monte carlo. CRC press.
Lao, J., Suter, C., Langmore, I., Chimisov, C., Saxena, A., Sountsov, P., Moore, D., Saurous, R. A., Hoffman, M. D., and Dillon, J. V. (2020). tfp. meme: Modern markov chain monte carlo tools built for modern hardware. arXiv preprint arXiv:2002.01184
Cabezas, A., Corenflos, A., Lao, J., Louf, R., Carnec, A., Chaudhari, K., Cohn-Gordon, R., Coullon, J., Deng, W., Duffield, S., et al. (2024). Blackjax: Com- posable bayesian inference in jax. arXiv preprint a7X1v:2402.10797.
Magnusson, M., Torgander, J., Bürkner, P.-C., Zhang, L., Carpenter, B., and Vehtari, A. (2025). ttetttt {posteriordb}: Testing, benchmarking and developing bayesian inference algorithms. In The 28th International Conference on Artificial Intelligence and Statistics.
Cabezas, A. and Nemeth, C. (2023). Transport elliptical slice sampling. In International Conference on Artificial Intelligence and Statistics, pages 3664-3676, PMLR.
Cechin, A. L., Regina, D., Simon, P., and Stertz, K. (2003). State automata extraction from recurrent neural nets using k-means and fuzzy clustering. In 23rd International Con- ference of the Chilean Computer Science Society, 2003. SCCC 2003. Proceedings., pages 73-78. IEEE.
Mira, et al. (2001). On cestropolis-haasting algorithms with delayed rejection. Metron, 59(3-4);231-22
Modi, C., Barnett, A., and Carpenter, B. (2024). Delayed rejection hamiltonian monte carlo for sampling multiscale distributions. Bayesian Analysis, 19535881-884.
DeRemer, F. L. (1969). Practical translators for LR (k) languages. PhD thesis, Massachusetts Institute of Tech- nology.
Murray, I., Adams, R., and MacKay, D. (2010). Elliptical slice sampling. In Proceedings of the thirteenth interna- tional conference on artificial intelligence and statistics, pages 541-548. JMLR Workshop and Conference Pro- ceedings.
Fan, J., Jiang, B., and Sun, Q. (2021). Hoeffding's inequality for general markov chains and its applications to statis- tical learning. Journal of Machine Learning Research, 22(139):1-35.
Muškardin, E., Aichernig, B. K., Pill, I., and Tappler, M. (2022). Learning finite state models from recurrent neural networks. In International Conference on Integrated Formal Methods, pages 229-2222222 29-248. Springer.
Flax (2023). Flax: A neural network library for jax. ht https: //tph/gttbb//////// com com/ggogle// flax. Accessed: 2025-01-30.
10

---

Efficienlly Vectorized MCMC on Modern Accelerators
Neal, R. M. (2003). Slice sampling. The annals of statistics, 31(3):70--767.
Zeng. Z.,, Goodman, R. M., and Smyth, P. (1993). Learn- ing finite state machines with self-clustering recurrent networks. Neural Computation, 5(6)9976-990.
Oktay, D., Mirramezani, M., Medina, E., and Adams, R. (2023). Neuromechanical autoencoders: Learning couple elastic and neural network nonlinearity. In Inter- national Conference on Learning Represenlations (ICLR 2023).
Pfau, D., Spencer, J. S., Matthews, A. G., and Foulkes, W. M. C. (2020). Ab initio solution of the many-electron schrödinger equation with deep neural networks. Physical review research, 2(3)033322..
Phan, D., Pradhan, N., and Jankowiak, M. (2019). Compos- able effects for flexible and accelerated probabilistic pro- gramming in numpyro. arXiv preprint arXiv:1912.1155:
Qiu, Y.. Zhang, L., and Wang, X. (2019). Unbiased con- trastive divergence algorithm for training energy-based latent variable models. International Conference on Learning Representations.
Radul, A., Patton, B.. Maclaurin, D., Hoffman, M., and A Saurous, R. (2020). Automatically batching control- intensive programs for modern accelerators. Proceedings of Machine Learning and Systems, 2:390-399.
Ruiz, F. J., Valera, I., Svensson, L., and Perez-Cruz, (2018). Infinite factorial finite state machine for blind multiuser channel estimation. IEEE Transactions on Cog- nitive Communications and Networking, 1(2):177-191.
Schoenholz, S. S. and Cubuk, E. D. (2021). Jax, md frame- work for differentiable physics. Journal of Statistical Mechanics: Theory and Experiment, 2021 2021(12):120116
Sountsov, P.. Carroll, C., and Hoffman, M. D. (2024). Run- ning markov chain monte carlo modern hardware and software. arXiv preprint
Svete, A. and Cotterell, (2023). Recurrent neural lan- guage models probabilistic finite-state automata. In Proceedings of the 2023 Conference on Empirical Meth- ods in Natural Language Processing, pages 8069-8086.
Tiño, P., Horne, B. G., Giles, C. L., and Collingwood, P. C. (1998). Finite state machines and recurrent neural net- autks-at and dynamical systems approaches. In Neural networks and pattern recognition, pages 171-219. Elsevier.
Vershynin, R. (2018). High-dimensional probability: An introduction with applications data science, volume 47. Cambridge university press.
Yeh, I.C. (2018). Real Estate
Valuation. DOI:
UCI Machine Learning Repository. http://oo...o//11.2242/////////

---

Efficiently Vectorized MCMC on Modern Accelerators
A. Mathematical Appendix
A.1. Proof of Theorem 4.1
Proof. The result broadly follows from known Hoeffding bounds for Markov chains. For clarity we restate the relevant result below", using our notation and set-up.
We first define formally the notion of absolute spectral gap, from (Fan al., 2021).
Definition A.1 (Absolute Spectral Gap). Let Z and denote the state space and the invariant measure of a Markov chain (Zihe1. For a function f Z + R. write (hh : h(2) (dz). Let (z(
{ [h: T(h?) < oo} be the Hilbert functions, and (=) = { h E (,2 ((h) 0) be its subspace of n-meanzero
space consisting of cu-square-integraan functions. The transition probability kernel of the Markov chain, denoted by P. is viewed as an operator acting on L2(x). is the absolute spectral gap of the Markov chain.
Let À E [0, 1] be the operator norm of P acting on (T). Then,
Proposition A.2 (Hoeffding Bound for Markov Chains Theorem in Fan et al. (2021)). Let (Zi)a1 be a Markov Chain with measurable state space Z, stationary distribution æ, and absolute spectral gap 1 X (0, 1). Let f 2 - [0, B) be measurable and bounded. Then, for any >> 0,
1 A 2c
Px Σ f(Z.) -BB,f(Z1) 2c
[ 2exp
(14)
1+ B
where 1 - is the absolute spectral gap of π.
Now we are ready to prove our results. We start with Co(m, n). First, note that since (X N1g)21 a Markov chain with stationary distribution æ for every chain j. then (fX,,, 1 1)221 is also a Markov chain with stationary distribution Th
w, because the chains are independent. Therefore, if we Z₃ (Xij, Nogl and
f(Z.) == maxje(m) (Ni,)= E [0, B), we get by an application of Proposition A.2 that
1- 26
Pa Σ max (Ni.) n., max (N15) 2 S exx jElm]
(15)
(I
_E(m]
Setting 5 = exxp
(-
1-> 1-1 2ca 1+1 n/2
and re-arranging, we get that with probability
5 (under the stationary distribution x),
n
2 a04
Σ
max (Ns.s) T.EE max (N1.)) jElm]
< B
In
(16)
jE ml
2(1
Multiplying by (m( and dividing by on both sides, and adding and subtracting Ag(m) from the LHS gives us the result for Co
1
(1+ X)
2
Bu (m)
Σ
max (No,) Bo(m) max max (N1.)) JElm)
$ Bo(m)B
In
(17)
n
JE(m]
2n A)
(1+ X)
(2)
Bo(m)
max (Ne,) + Ao(m) Bo(m) max (N %Elm]
$ Bo(m)
In
(18)
12
JElm
2n(1 - A)
(1+A) In
2
Ca(m, Ao(m) Ex max (N (N
< B
(19)
jE[m]
2n(1 X) 8
'We note that Fan al. (2021) only present one-sided bound, but standard symmetry arguments this immediately implies the above two-sided bound.
12

---

Efficientlyy Vectorized MCMC on Modern Accelerators
Now follow similar steps for Cp(m, terms of sum of individual distances using the union bound.
n). n). To start, bound the distance from maxjem 1 n Ex=1 Mi and Ex [N11 in
Pr
max jem
Ni. - Nn 2e = Px
max jem 1
Σ
Ni. - NN 20)
+ Ps
max JEM
Nis - EN S
(20)
No. - N,N E 20) ))
Px
j=1
Ni. - E, S -00)
+
(21)
m
1
$
Σ
P
No ENNU1
11
j=1
1
+P,
NJ. NI
(22)
1/1
P
No. - E Nu1 ZE
(23)
Σ NE.1 NN Ze
mPr
(24)
mP,
Σ
No.1 nÉ N11
N nE
(25)
Applying Proposition A.2 on the Markov Chain (X N1,1)12 with f(X,,1, N1.1) NI. E [0, B] and following the same steps as for Co(m,n), n), we similarly recover
(1+ In 2n [ X) 8
2m
CF(m,n) AF(m) Br(m) max (Ni,)) jE[m)
< BB(m)BB
(26)
which is the result in the Theorem.
Proposition A.3. Fix m, K E NY10) and let PN be probability measure on R4, strictly positive first moment. Suppose (i)
N1,..., Nm ~ PN. (ii) (m()) Cx (m) 2 0 and (iii) E [maxxe{[]] (mm) EjE[k] (,(m), 1). Then, we have
(m) + ( Nj E(m) a(c=((m) + (K)) - EN,)
E max Nj R(m)
(27)
EN1
where e (m)) = j C (m). The bound is tight.
Proof. Note a+d ctd bn + (1 - 7) where 7 = aid =i for any a, b,b, R. Applying this to our case, we get
(m)
(m(m)
E(m) a(c= (m) + (m()))(( 1) w + a(c=)(m) + cm()))
R(m)(1 - w)
(28)
a(c= where afe ( E [0, 1]. Now we split into two cases for K = and K > 1. For the case K = 1 we only have a single iterative state and so C (m) = 0, a 1. In this case we trivially have F(() R(m). Now
13

---

Efficiently Vectorized MCMC on Modern Accelerators
suppose K 1. In this case, since (m)- + cx(m)) 2 maxje[K] (m)), have
(m(m)
cx(m)
[a (m) S 1
(29)
((c=x(m) + ca (m))
Which means we can bound Emm) by removing the term in front of R(m),
(x(m)
[(m) S alc-n (m) + (x(m))(K 1) w + (1 -ww w)
(30)
By the same logic, we have
(m)
(m)
Ejnecsln
S
maxje[K] (m) - S
(31)
a(c-x (m) + (m))(( 11)
maxje[k] ( -
which means
E(m) S w + R(m) - w) I R(m)
(32)
(33)
Where the last line uses the fact that R(m) E(m) = R(m). This completes the proof
2 1 since N₁ 2 0. The bound is tight because when K = we have
14

---

Efficiently Vectorized MCMC on Modern Accelerators
B. Additional Details
B.1. Algorithms
Note here we use to denote a batch inputs [T1, Tm.] for m different chains, and the same for other variables.
Algorithm Vectorized MCMC algorithm with vmap (sample) function
1: Inputs: sample øo. seed Fo
2: for E {1...., n do
3: generate 0i, f- vmap (sample) (0--1,77-1)
4: end for
5: return 11....,0,, S
Algorithm Vectorized FSM MCMC algorithm with vmap (step) function
1: input: initial value To, # samples n
2: initialize: = vmap (init) (ão), X listO, B list()
3: S 0 and h 0
4: while min;, [ < n do
5: (k: it.,imsamppee) + vmap (step) (k, 8)
6: append current sample value & stored to X
7: append sSample le to B
update sample counter f issample
9: end while
10: return [[B]
Algorithm 8 Transition kernel for elliptical slice sampler with log-pdf log covariance matrix E.
1: Input: Sample
2: Choose ellipse ~ ~O, ))
3: Set threshold log + log p(x) + log u ~ 10
4: Set bracket [0min mmxx] + 10 2n, 0] ~ 21 2m]
5: Make proposal x' + x cos + v sin 0
6: while log p(x') > log y do
7: Shrink bracket and update proposal:
8: if < then
9: 0min +- 0
10: else
11:
0mnx +-
12:
end if
x' f- x cos + sin 8: [1(0min max]
end while
15: Return x
Algorithm 9 Soundled.steg as input to amortized step.
1: Input: Algorithm state kt variables
2: doComoutation = False
3: while not doComputation do
4: (ks, 2, i isampll sssaap/le..ccmmpptcccooo//a/s/// step(k,
5: end while
6: Return 2,, is Sammlle, doComputation)
15

---

Efficiently Vectorized MCMC on Modern Accelerators
B.2. FSM Construction Procedure for Programs Considered in Main Text
Detailed Construction in the two sequential while loops case Here, we describe the FSM construction for programs with two sequential while loops in detail. Following the constructions introduced in the main body, let us note
Fi = ((S1) S12, S13), 61, S1a) the FSM associated to Bi
F ([S21, S22, S23), 2, 02, S21, S2a) the FSM associated to I32
By construction, S21 is empty. Both FSMs share the same input space, which the set of local variables values associated the original sample function.
Then, the resulting FSM representation of sample is 2, S23), where
S (S11, S13, S22, $23)
Transition function á defined
if (S1,, S12) if S S13
=(S,
82(221
(34)
(2(
if S = (2,, S23
whose construction illustrates the "FSM stitching" operation performed.
Detailed Construction in the two nested while loop Here, describe the FSM construction for programs with two nested while loops. Following the constructions introduced in the main body, let us note
F, = ([Sh, S2, S.a) Z, 6,, S11, Sta) the (inner) FSM associated to B2.
Then, the resulting FSM representation of sample (So, Z, 8 S1, Sz), where
S. = [S1, Sn. S22, Sa. Ss}.
Transition function 5, defined as
6S,, =) = 6.( 2) runs the outer while loop condition goes to Sn1 if True, and S3 otherwise. -8, 6,(8, 2) S Sp2
B.3. FSM Construction Procedure for General Programs.
The FSM construction procedures in ?? for specific programs can be automated to handle programs with any finite number of sequential and/or nested while loops. We show how this works using simple programming language defined by the following grammar:
Operation op ==
op
op
op...
Simple Block SB
(op)
(op)
(SB)
While Block WB
W (Bs)
Block B :
(SB) (WB)
Blocks Bs ===
(B) (Bs) (B)
This grammar contains just enough to describe non-trivial programs containing any finite number of (possibly nested) while
loops. Given program generated with this grammar, we construct an FSM variant by:
1. Parsing the program into a parse tree.
2. "Coarsening" the parse tree into a tree where each node is block.
3. Applying recursively-defined function which transforms the program into an FSM.
4. Collapsing any spurious empty states.
16

---

Efficiently Vectorized MCMC on Modern Accelerators
B.3.1. STEP 1: PARSING THE PROGRAM INTO A PARSE TREE.
First, note that as this grammar is context-free, every program P can be represented by the yield of parse tree (e.g. Theorem 5.12 (Hopcroft et al., 2001)). Consequently, choose the look- aook-a left-to-right, rightmost derivation (LALR) parser (DeRemer, 1969). Figure 8 shows the resulting parse tree associated with the program "opl (opl W {""
use any context-free parser to obtain parse tree of the program. We
Farry Wy
Tewer
Find BUA
Figure 8. Parse tree of the program "opl (opp:o//1W(()))"
B.3.2. STEP 2: COARSENING THE PARSE TREE.
To convert the parse tree into an FSM, we first derive a simplified version of with: (i) no operation nodes (c.g. content of simple blocks), (ii) no block markers (e.g. While tokens and braces), (iii) no "blocks" and "block" nodes (these are helpful to clarify the grammar's 's definition. but do not carry information about a given program). The first two properties can be achieved by simply pruning the tree of all nodes that are not blocks, which is well-known operation and not shown here. The output of this procedure is guaranteed to be tree, as all such nodes (i) and (ii) are terminal. The second property can be achieved by applying the procedure in Algorithm 11.
Algorithm 10 Gathering meaningful children
1: function GATHER_CHILDERR CHILDREN((Odd returns list of Node
2: all children U MAP(GATHER CHILDREN, chilechildren)
3: all_children U COTFER(NOOTTEMMINAA all acl_dchirr)n) (Discard "op" nodes and markers)
4: if node.type is "block" or "blocks" then
5: return all children (Discard the node, and only return the children)
6: else
7:
return (Node(node.type, all [cildren)]
8: end if
9: end function
Algorithm 11 Coarsening the parse tree
function COOARSEN_TREE/ON returns Node
2: all children U MAP(GATHER CHILDREN, chiddren(node))
3: return Node("blocks", clllchilr)
4:
end function
After applying the coarsening procedure to the parse tree of the same program (oplW(opIWoop1)) we obtain the coarsened parse tree shown in Figure 9. The algorithm has been modified to return block labels in order to track the graph manipulation operations done in the next steps.
B.3.3. STEP 3: TRANSFORMING THE COARSENED PARSE TREE INTO AN FSM.
Now we transform the coarsened parse tree into an FSM. Our approach leverages the fact that any subtree of the parse tree itself valid (sub)program and should thus be handle-able by the conversion procedure we define. Each (sub)program
't is "suhprogram" of the original program sense that it is contained in the original program.
17

---

Efficiently Vectorized MCMC on Modern Accelerators
simple_bloce
simple_block2
blocks
1alle_block_1
chile_block 2
simple _bl__blook 3
@mple_block.
@mple_block__ _
Figure Coarsened parse tree of the program "op]W[opl[(<pp))""
either a single simple block, an arbitrary while block (looping over a sequence of blocks), or the root program (c.g. a sequence blocks).
This suggests recursive conversion procedure which computes single-state FSM for the block the leaf nodes, and then (recursively) combines these FSMs into larger FSMs at each layer of the parse tree, until reach the root node, at which point the final FSM is returned. The FSM at leaf nodes (i.e. simple blocks) is simple single state FSM. The FSM at non-leaf nodes (e.g. while blocks or the entire program) is constructed by (i) connecting the terminal state of the hth child's s FSM to the initial state of the (i + 1)dt child's FSM, for every child i of the non-leaf node except the last, and (ii) the node a while block, (a) adding an edge from the last child's FSM to the first child's FSM, and (b) assigning the initial and terminal state of the current node's FSM empty states with an edge between them. Additions (a) and (b) account for the fact that unlike standard blocks, while blocks have bodies that may be looped over, skipped.
Using empty states initial and terminal FSM of while blocks allows our FSM conversion procedure to be "context-free". each subFSM is obtained by assembling sub-subFSMs, without knowledge about the nature of the blocks being connected. Without these empty states, we would require more complicated logic to perform the connection. For instance, in program of the form "opl W (opl) W(opl oop2", one needs to connect the first op1 to the last op2 to encode the fact that both while loops may be skipped.
The full algorithm given Algorithm 12. The resulting FSM for the program "opIW W [opl}}" is shown in Figure 10. The nodes starting with an "T" (resp. "T"" the (empty) initial (resp. terminal) states of each while block. The transitions are labeled using the following convention:
"E(n)" means "enter loop number n²*
"C(n)" means "continue loop number n²
(n)))" means "skip loop number n²
"F(n)" means "exit loop number n²"
"N" means "next block" (which are created when connecting the different subFSMs of the root program)
B.3.4. STEP 4: COLLAPSING THE EMPTY STATES.
The empty states produced by the step above are side effects of the conversion procedure used: the final FSM should not contain them. To this end, we next describe a post-processing, iterative procedure which removes them from the FSM. At each iteration, the procedure selects (if any) an empty state, viddunless transitions into itself (i)revvess from the list of FSM nodes and (ii) adds edges connecting the parents of the empty state to children.
18

---

Efficiently Vectorized MCMC on Modern Accelerators
S1
$2
12
S3
T2
S4
T1
51
Figure 10. Intermediate FSM obtained from the coarsened parse tree of the program "oppW[oplWWopp] This intermediate FSM contains empty states, which are collapsed obtain the final FSM.
Algorithm 12 FSM creation from a coarsened parse tree
I: function COARSE. TREE TO_ FSM(node) returns FSM
2: if node.type is "while.blocks or "blocks" then
3: C. ffsms + MAP(COARSE. TREE. TO..TOOFFSS.. FSM, node.children) f.fmm))
4: nodes + U MAP(GET
5: edges + MAP(GET EDGES, fisms)
6: labels U MAP(GET_LABEI, c_fsms)
7: // Connect the different subFSMs
8: conn sonn_edees {(c_ (oc.ffms(-l].eermnall c. finiti). for = 1,... Len(c_fsmss 1}
9: comn_labels - ("N",... "N"}
10: edges, labels < edges U conn_edges, labels cann_labels
11:
it- COUNTER(
if node.type == "while block" then
terminal_state, initial_state Node("terminal"), Node("initial")
14:
enter_loop_ede (initial_state, c_fisms(0).nitiaal)
15:
exit_loop_edge +- (c_fsma(--] terminal_star _stittel__stt))
16:
continue loop loop.edge + (c. C.Cffms[-]].ermmnndd с fsmms(0).nnitia)
17:
hhtplloop_edge < (terminal state, initial state)
18:
nodes <- nodes [terminal-state, initial state]
state)
edges < edges U Ccontinue_loop.c "C" +i, "S" +i}
Akip_lop_eed .dgee, enter_loop_edge exit_ edge)
20:
labels labels ("E"+i, "F"" +i, "C"
21:
return FSM(nodes, edges, labels)
22: else
23:
return PSM[[simpee.locc// (].(0)
24:
end if
25:
end if
26:
end function
19

---

Efficiently Vectorized MCMC on Modern Accelerators
The labels of the new edges obtained by concatenating the "parent-to-empty state" and "empty state-to-children" edges. The procedure is given in Algorithm 13, and the resulting FSM for the program "opl W ( ) is shown in 11. The procedure results in FSMs that agree those in the main text for single while loops, two nested while loops, and two sequential while loops. There are two minor caveats, which we discuss below.
Figure
Remark B.1 (Handling empty states with self-transitions). The procedure above removes all empty states, apart from those with self-transitions. These self- trrns-rrnniiinssssoo not present before collapsing the empty states, and arise in the midst of the procedure, after collapsing subset out of all the FSM empty state. Importantly, such self- self-transitions are pathological: the state of the program does not change when performing a self-transition into empty node, if such self- oeffttransition occurs, it must keep occurring indefinitely afterwards. Thus, programs resulting into FSM with self-recurring empty states form class of "potentially non-halting" programs.
Remark B.2 (Handling impossible transitions). The node collapsing process results in the creation of "multistep" edges. These transitions involve evaluating multiple conditions in single step. Sometimes, these conditions mutually exclusive: in that case, this multistep transition is impossible, and can be removed from the graph. By removing such transitions after each step of the procedure, the routine can be shown to not add any duplicated edges to the FSM, as proved in the next proposition.
Algorithm 13 Collapsing empty states
function COLLAPSE EEELLPPPEEEETY...M EWW.EEPPYYYE ((M)T)))TTES(((M)) returns FSM
2: edges <- fim.edges
3: labels +- fsmllabels
4: for all node in FELPER NOT INITIAL NoT FINAL, nodes) do
5: if node E node.children then
6:
continue
7: end if
8: for all parent, child in PRODUCT(node.parentt chieechildren) do
9: new edge (parent, child)
10:
if new edde edges then
11:
error (Edge already exists}
12: end if
13:
new label < AND(parent.aabel laild.label)
14:
IMPOSSIBLEEnewLLae label) then continue
16: end if
17:
edges edges U (new_ed (ewee
18:
labels <- labels J (new_label)
19:
end for
20:
end for
21:
nodes + FILTER(NON JEPPT,, fsm.nodes)
return MSodds,ss, edges, labels)
end function
20

---

Efficiently Vectorized MCMC on Modern Accelerators
C2
(N + El)
(N + E2)
(F2+N)
(FI+N) N)
S1
S2
S3
S4
S5
S2
+N)
Figure 11. FSM for the program "opl W [opl [op]W(op1))) with empty states collapsed, using the automatic procedure. Note that this matches the structure of the FSM for the No-U-Turn sampler in Figure 2, which was produced using (essentially) a special case of this procedure programs with two nested while loops.
21

---

Efficiently Vectorized MCMC on Modern Accelerators
B.4. Implementation Details
FSM Step Function Design. At each step of the FSM, one always executes block S4 s' which updates the inputs, before calling the transition function 8 on (k, Sx(z) to determine which block should be next. One therefore combine these into single (state) function T () (o(k, (k())), ( Sz(2)), equivalently, collection of functions Ti...., Tk where T3 + (ō(k, (z)),, Sx(=).. In practice, our implementation of Algorithm 2 calls single swit, th over these composite state functions to streamline the code. The same composite state functions are also used in our implementation of bundled step in Algorithm 4.
FSM Runtime Design. In our experiments, we use native Python while loop in Algorithm (and its vmap-ed variant (JIT compiled) jax lax. scan run the FSM step function for blocks of t 100 steps, when drawing 100 samples. This gives the flexibility store the results in dynamically shaped lists/arrays and transport to the CPU for faster array slicing when CPU memory is available, whilst still reaping the benefits of JIT compilation.
in Algorithm 7). Inside the loop,
Compilation. We JIT compile both vmap (step) and vmap (s ammple) functions for each MCMC algorithm implementa- tion (here step refers to any of the basic variant Algorithm 2, bundled variant Algorithm 4, amortized variant Algorithm or the bundled and amortized variant Algorithm 9). For Delayed Rejection and the Elliptical Slice Sampler (with n 25) remove compilation time get accurate results the runs with small numbers of chains m, due the low cost of the computations involved.
JAX implementation. When comparing to non-FSM implementations, we used BlackJAX (Cabezas et al., 2024) for fair comparison with our method, since we use BlackJAX primitives for the key computations in some of our algorithms (c.g. HMC-NUTS). Where BlackJAX implementation was not available (e.g. Delayed Rejection), we wrote our own for fair comparison with our FSM implementation.
Bundling with Amortized Step. As discussed in Section 5, one cannot non-trivially use bundled_sted as the step function inside amorti poorrized step, because given sequence of states may require the amortized computation to be updated in the middle, and nothing in bundled_ste amortizing an expensive computation g (e.g. expensive log-pdf calculations), we use another step function defined in Algorithm 9, as the 'step' function called inside flag) until doComputtion
flags this. To reap some of the benefits of step bundling when
This step function iteratively runs step (modified to True ironically, using a while loop. When vmap-ed, each
return the doComputat 'step' inside amort wed__tep now executes sequences of cheap states that not require g. until g is required again for all chains. This works effectively when g called inside the iterative state, which is typically the case when g = log p.
B.5. Additional Results
n 25
n 100
n 400
2
Efficiency Ratio
21
21
20
20
iters per sample walltime
20
20 2
2
23
24 # chains
25
27
2
20
20 2
2
23
24
25
26
20
2
22
23
24
25
26
27
2
29
20
# chains
# chains
Figure 12. Efficiency Ratio of our elliptical slice FSM against BlackJAX's elliptical slice algorithm (as measured by estimated R(m) [[maxyejm] N,]/B[N,] (i.e. iters per sample) and walltime) the Real Estate Dataset described in Section 7 when restricting the dataset to the first n E (25, 100, 100) datapoints. The relative efficiency of the FSM improves as the number of chains used increase, and the log- likelihood increases. When 400, almost achieve the theoretical bound R(m) in speed-ups.
22

---

Efficiently Vectorized MCMC on Modern Accelerators
n
25
n
100
400
27
210
24
Walltime
26
28
25
2
24
26
20
21
22
2
24
2
2
2
2
29
20
20
21
2
2
2
2
2
2
2
20
2
2-
2
2
24
2
2
27
2
2
217
FSM slice
ESS per second
255
BlackjAX slice
213
2
233
21
2
211
29
29
26
27
20
2
2
2
24
25
26
27
2
2
200
29
2
2
2
24
25
26
27
2
29
20
2
2-
2
2
24
2
2
2
2
29 20
chains
chains
# chains
Figure 13. Walltimes and ESS per second using the Elliptical Slice Sampler (non-FSM vs FSM implementation) on the Real Estate Dataset described in Section 7, when restricting the dataset the first E {25, 100, 400) datapoints. For each dataset size, the best walltime and ESS/sccond is obtained by both implementations when using m = 1024 chains. Our FSM implementation can obtain the greatest efficiency for all dataset sizes. the lieliiol cost increases (the log-likelihood in GPR regression costs Onn the FSM efficiency gain increase, reflecting the benefits of amortization.
Distribution
NUTS-FSM
Standard NUTS
FSM Speed-up
Real-Estate GPR
863.6
274.1
3.15x
Soil
0.014
0.006
2.43x
Pilots
3.538
3.869
0.91x
Table ESS/see comparison NUTS with without FSM acceleration on Real-Estate GPR sampling problem considered Section 7.2, and two sampling problems (Soil and Pilots) from the PosteriorDB database (Magnusson al., 2025). Results averaged over 1000 samples and 128 chains. As with Section 7.4 in the main text, we use 400 warm-up steps to tune the mass matrix and step size of FSM implementation, and the BlackJAX (non-FSM) implementation. We find substantial speed-ups in 2/3 problems. Both of these problems (Real-estate GPR and Soil) have an expensive log-pdf (and gradient) whilst the log-pdf in Pilots is much cheaper. In general, when the number of integration steps taken NUTS is moderate high on average (so that most time spent integrating along the trajectory), an expensive log-pdf makes the additional cost of executing all other blocks in the FSM step function (when integrating) relatively smaller. That is, the gap between (mm) and R(m) (as defined in Scction is smaller, hence increasing the relative efficiency of the FSM. The opposite happens when the log-pdf is cheap (as in Pilots). We also note that computational resources did not balance out across chains Soil and Pilots after 1000 samples. therefore would expect see larger speed-ps when running for longer.
23


Reviews:
[{"review_id": "doIPfJZ23X", "reviewer": "Reviewer_akMN", "summary": "This work proposes a way of more effectively parallelising (multiple chains of) certain MCMC algorithms via automatic vectorisation tools like the vmap function in JAX. Specifically, this work is concerned with MCMC algorithms whose computational cost per iteration has a large variance, e.g., due to the use of a proposal based around rejection sampling. \n\nThe authors show that a naive parallelisation of such chains (i.e. by passing functions that execute an entire MCMC iteration to vmap) is wasteful because it leaves most of the processors idle while waiting for the slowest iteration to complete. Instead, they propose to break MCMC iterations down into smaller steps (formalised via finite state machines). This allows automatic vectorisation to occur on the level of these smaller steps rather than on the level of entire MCMC iterations.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review and recommendations. We are glad that you see the value of our contribution and your feedback has been helpful in strengthening the clarity of the paper. Below we address your main questions/comments:\n\n**On step bundling**:  In short, performance should always be weakly improved (i.e. not get worse), since at worst `bundled_step` takes a single step during each call. \n\n*In more detail*: \n- As long as the switch in `vmap(step)` executes the blocks/states sequentially (and our testing indicates that it does), the step-bundling routine should always weakly outperform the basic `step` function. Like `vmap(step)`, `vmap(bundled_step)` will execute all blocks/states (sequentially) when called. However, `bundled_step` allows any chains that transition from $S_{k} \\to S_{k+1}$, to (for example) immediately transition from $S_{k+1} \\to S_{k+2}$ at no extra cost (assuming $S_{k+1}$ is called after $S_k$  in the step bundling order). As such, in the worst case each chain only progresses through exactly one transition (which mirrors the behaviour of `step`).\n- Whilst one can optimise the ordering of `bundled_step` to improve performance, a surprisingly effective heuristic is just to use the chronological order of $S_1,...,S_K$, since (for all algorithms we considered in this work) $S_{k+1}$ is constructed as a block that is callable after $S_k$ in the original `sample` function. This ordering is what we use in the paper. In our experiments, we have not found a significant gain to manual tuning of the block ordering over chronological ordering.\n\n**Definition of $c_{\\neg k}$**: Thank you for pointing out this omission, we meant to define $c_{\\neg k}(m) = \\sum_{j \\neq k} c_j(m)$.\n\n**Eq(10)**: Equation 10 holds asymptotically as $n \\to \\infty$, as it is derived by plugging in the limits of $C_0(m,n)$ and $C_F(m,n)$, which are given by Theorem 4.1. Since the approximation error for each of these terms is (with a high probability) $\\mathcal O(1/\\sqrt{n})$, after some basic manipulations (and assuming $C_F(m,n) > 0$) this results in the approximation error $C_0(m,n)/C_F(m,n) - E(m) = \\mathcal O(1/\\sqrt{n})$ (again, with some fixed high probability $1-\\delta$). If space permits and there are no objections, we are happy to add this analysis in the main text.\n\n**Other typos+ formatting suggestions**: Thank you for pointing out the typos, formatting and improvements, we will amend those accordingly in the final version.", "claims_and_evidence": "The claims are supported by clear and convincing evidence. The computational cost improvement over \"naive\" vectorisation (on the MCMC iteration level) is established theoretically. Run-time improvements are demonstrated empirically on a number of benchmark examples.", "methods_and_evaluation_criteria": "The proposed methods make sense.", "theoretical_claims": "I checked the main manuscript but not the additional derivations in the appendix. I did not find any issues.", "experimental_designs_or_analyses": "I only checked the experiments insofar as checking their description in the manuscript. I did not find any isues.", "relation_to_broader_scientific_literature": "Effectively using parallel architectures to speed up MCMC implementations is an important research area. This work contributes to it by providing a fairly simple solution to exploit existing general-purpose automatic vectorisation routines.", "essential_references_not_discussed": "None come to mind.", "other_strengths_and_weaknesses": "I think this work is overall well written and well structured. As far as I know, the contribution is original. In my opinion, the contribution is also significant enough for publication in ICML because it provides a fairly simple way of speeding up multiple-chain MCMC sampling (for certain MCMC kernels) using existing general-purpose automatic vectorisation functionality.", "other_comments_or_suggestions": "- Figures should not appear in the text on pages before they have been referenced (see, e.g., Fig. 1 which is only referenced on Page 3; or Fig. 6)\n- Line 204, RHS: the double use of the index $k$ in the lower bound for $\\alpha$ is confusing\n- Section 4, 2nd paragraph. Maybe denote the index of the block containing the loop by $k^*$ rather than $k$ to avoid confusing this \"special\" value of $k$ and $k$ as a generic index below\n- Line 2020: define what $c$ with \"negated $k$\" in the subscript means\n- Eq. 10: Is this really equal for finite $n$ (not just approximately equal)?\n- Line 284: use used\n- Fig. 7: \"and or than half\" in the caption\n- Table 1: use natbib's citet rather than citep for the reference in the caption\n- Bibliography: fix missing capital letters in names/proper nouns in article titles, and in journal names; use of journal name abbreviations is also not consistent.", "questions_for_authors": "Can you give some more guidance for how to optimise the \"bundling\" for the \"bundled_step\" routine? Presumably, one has to be careful here because its performance can end up being worse than \"step\" (i.e. without any bundling).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1rxoluiRUR", "reviewer": "Reviewer_pDxh", "summary": "Vectorizing Markov Chain Monte Carlo (MCMC) algorithms using vmap to create multi chain code results in a synchronization problem. That is all chains have to wait for the last chain to be completed. The paper proposes to solve this problem by using FSMs, finite state machines, to design single-chain MCMC algorithms in a way that can avoid synchronization barriers when vectorizing with vmap. The authors implemented popular MCMC algorithms like slice sampling, hmc-nuts, delayed rejection.It is possible to write single chain MCMC code and call vmap to turn it into vectorized, multi chain code that can run in parallel on the same processor. Vmap transforms every instruction in the sample into a corresponding instruction operating on a batch of inputs matrix-vector multiplication to vectorize the sample. These instructions are executed in lock-step across all chains. Those MCMC libraries adopted ML frameworks with automatic vectorization tools as their backend. Bottlenecks for these libraries happen at control flow, if there is a while loop all chains wait for the last chain. Paper is showing transforming MCMC algorithms into equivalent sampler that is avoiding mentioned synchronization problem; the novelty is at the approach of transforming MCMC algorithms into FSMs.These FSMs enable us to breakdown ‘sample’ into series of smaller steps that has minimal variance in execution time and avoids using while loops.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our work and for your recommendations to improve the paper. Below we address your comments and questions. \n\n**More general while loop structures:**  The conversion of algorithms into FSMs by splitting up while loops is a generally applicable recipe, but we agree that our presentation in Sec 3 sweeps this under the rug by immediately launching into specific algorithms. We have added a paragraph at the outset that explains how to obtain FSMs for any finite number of while loops. This can in fact be done automatically: We have a parser that reads the MCMC algorithm in a symbolic language and outputs the FSM (see below for details). This is already implemented and tested. We had not included it to save space and because all MCMC algorithms with while loops we are aware off fall into the categories considered in our paper. However, since the automatic procedure exactly addresses your point, we will add it to the appendix if the reviewers have no objections.\n\n*General FSM construction procedure:* We have derived an algorithm able to convert arbitrary programs (written in a symbolic programming language) into their respective FSM graph. The algorithm (1) calls a parser to obtain the syntax tree of its input program, (2)  creates a coarsened version of that tree where each leaf corresponds to a code block (and each non-terminal node corresponds to a while loop condition), and (3) recursively transforms that tree into an FSM graph. Each node and edge of the final FSM graph can be mapped back to a code block or a while loop condition of the original program. \n\n**Step function errors:** Our method transforms an existing MCMC program into an equivalent sampler using the same lines of code within blocks. As such, any error handling in the original program is retained by the FSM. To handle preemption, one can straightforwardly checkpoint the state of the FSM and all flags at a desired frequency of `step` updates.\n\n**Bundling method:** In Section 5.1, we provide an illustrative example of `bundled_step` (see Algorithm 4) and discuss how it improves performance over `step`. We also tested the performance of `bundled_step` in the Delayed-Rejection experiment 7.1 (see the discussion under Experimental Setup, Results, and in the caption of Fig 6). However, we agree that more explanation of the effect of step-bundling would improve the paper. To that end, we propose to add text similar to the following in Section 5.1: \n\n\"As long as the switch in `step` executes each state/block sequentially when using `vmap` (and our testing finds that it does$^*$), then, both `vmap(step)` and `vmap(bundled_step)` have the same execution cost. This is because `vmap(bundled_step)` always executes all blocks/states sequentially, and both `vmap(bundled_step)` and `vmap(step)` execute all blocks. As a result,  if each chain makes on average M transitions when calling `vmap(bundled_step)`, the speed-up over using `vmap(step)` is M-fold. Since $M \\geq 1$ (i.e. `bundled_step` must make at least one transition each time it is called), this means that `bundled_step` must always weakly improve the performance of `step`, with `vmap`\".\n\n$^*$We have tested the JAX's `switch` and found that it behaves consistently with sequential block execution. We can add this to the appendix, if desired.\n\n**Amortization method:** We explain in Section 5.2 that amortization ensures any function $g$ that is called in multiple states/blocks, will only be executed once per FSM step. This reduces $\\alpha$. In Section 7.2 we also test the effect of amortization on the performance of the elliptical slice sampler (see line 376 onwards).  However, we appreciate that more detail on the effect of amortization would improve the clarity of our paper. We therefore propose to add text similar to the following in Section 5.1:\n\n\"If a function $g$ is called in M different states/blocks and its executions cost (in total) $\\beta \\in [0,1]$ fraction of the cost of all blocks/states  $S_1,...,S_K$, then amortization will reduce the cost of calling `vmap(step)` to $\\beta/M + (1-\\beta)$ fraction of the original cost, since $g$ is now executed only once, instead of $M$ times. We indeed observe this for the elliptical slice sampler in Section 7.1 (see Fig 5 - here $M=2$ and $\\beta \\approx 1$ as the log-likelihood is very expensive, resulting in a 2x speed-up over no amortization)\". \n\n**JAX vs. Torch/TensorFlow:** We agree that testing other frameworks is of interest. That we have not done so yet is because (1) PyTorch is not applicable, since its version of `vmap`does not yet handle data-dependent control-flow, and (2) TensorFlow’s `vectorized_map` is so similar to JAX’s `vmap` (it is a tensor-level vectorization tool that constructs a single batched while loop for all chains) that we expect it to behave very similarly. We do plan to test this properly in the future, but with limited space, we feel that a proper evaluation for one framework is more helpful.", "claims_and_evidence": "The claims in the paper are supported by experimental evaluations.", "methods_and_evaluation_criteria": "The proposed methods make sense for the targeted problem.", "theoretical_claims": "No theoretical claims in the paper.", "experimental_designs_or_analyses": "The experimental analysis makes sense but it is not possible to validate the results.", "relation_to_broader_scientific_literature": "-", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "Strengths:\n- Sample to FSM approach is generalizable for various algorithms\n-Time complexity of FSM and corresponding functions are given; comprehensive detailed mathematical formula and derivation is given in detail\n-Various MCMC algorithms are used including Delayed-Rejection Metropolis Hasting, Elliptical Slice Sampling,HMC-NUTS their approach is tested on various different algorithms\n\nWeaknesses and Suggestions:\n-The paper shows FSM of an MCMC algorithm with a single while loop, two sequential while loops, two nested while-loops. It does not mention different configurations and how FSM should be obtained if these while loops occur in different way or interacting each other\n-Step function is defined, this function performs single transition along an edge in FSM diagram, in definition of step function it is called iteratively over blocks of samples, this process given simply not mentioned how to deal with errors, or where to resume if long sequence of blocks of sample preempted for some reason\n-It shows how to make FSM design more optimal, by reducing parameters α and K, Paper claims step bundling method increases efficiency, but no detailed example is given\n-Another optimization is proposed as returning additional g_flag variable for step function; but, the affect of it on cost function or time complexity is not discussed\n\nMinor:\nSome experiment results are shared on jax. Although jax is popular, pytorch or tensorflow could be used to extend derivation of their approach in different frameworks and libraries", "other_comments_or_suggestions": "none", "questions_for_authors": "see the points listed under \"weaknesses\" above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Cn17HkJkWX", "reviewer": "Reviewer_UagZ", "summary": "This work focuses on minimizing synchronization steps when vectorizing MCMC algorithms. Modern MCMC methods, such as NUTS, involve stochastic number of operations per Markovian transition, depending on the initial sample and random seed. This introduces unavoidable overhead when running multiple chains in a vectorized manner, as each transition step requires synchronization, forcing all parallel chains to wait for the slowest one to complete, leading to wasted computational resources.\n\nAddressing this long-standing challenge is crucial for enabling MCMC algorithms to fully leverage modern hardware, such as GPUs. Most existing approaches tackle this problem by designing new MCMC algorithms that minimize random execution steps (e.g., CHEES). In contrast, this work provides a systematic implementation strategy for existing algorithms—such as NUTS and the slice sampler—using finite state machines on a single chain, ensuring that synchronization occurs only at the final transition step during vectorization.\n\nThe authors present theoretical results quantifying the expected performance gains compared to naïve vectorization and empirically demonstrate the effectiveness of their approach across multiple examples.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and engaging review, we are excited that you see the value in our contribution. Below we address your comments and questions.\n\n**FSM Efficiency gain:** The main reasons that the improvement in walltimes don’t match iters/sample in Fig 6 are (i) as you guessed, the slowest chain is still substantially slower than the mean after n=1000 samples, and (ii) the log-likelihood is very cheap to compute in this experiment, so control-flow costs (which are large in NUTS) dominate. Since each FSM `step` executes all states (see clarification below on branch evaluation cost vs. naive vectorization), this makes the trajectory integration more expensive per step, even with step bundling. It should also be noted that $R(m)$ is around $20$ for $m=100$ chains, and we obtain a speed-up of around 10x, so the gap is actually not that large. For $m=500$ chains it is larger, but we suspect this is because of limited GPU capacity.\n\n**FSM evaluating all branches:** All branches of control-flows are indeed evaluated when vectorizing both the FSM and the naive implementation with vmap. However, this has different effects in each case due to the differing control-flow used “between” the blocks/states. The FSM `step` uses a switch over all code blocks $S_1,…,S_K$ that comprise `sample`. This means that a single call of `vmap(step)` has to execute all blocks for all chains, whichever block the chains are currently on. By contrast, in `sample`, the control-flow “between blocks” takes the form of (possibly multiple) while loops. In the case of a single while loop with blocks $S_1,S_2,S_3$, this results in the loop body ($S_2$) being evaluated for all chains until they have finished iterating. However when calling $S_1$ and $S_3$ with sample, only those blocks are executed. \n\nEssentially, we exchange the variance in execution time in `sample` (across chains), for a bias. As a result, when there is little/no variance in execution time (e.g. the while loop always terminates after 1 iteration), the basic FSM (without our optimizations) can perform worse than naive vectorization. The step-bundling and amortization procedures outlined in Section 5 are specifically designed to minimize this bias. For example, by using `bundled_step` (with blocks called in chronological order) the FSM would run in the same time as the naive vectorization if there is no waiting and the loop terminates in one iteration, because a single step now progresses all chains through $S_1 \\to S_3$. Amortization enables functions which are executed in multiple blocks/states to only be called once per FSM step, therefore reducing the additional overhead. For the elliptical slice sampler, doing this for the logpdf can reduce the cost of executing all blocks almost back to the cost of a single block, whenever the logpdf is expensive enough to dominate computational cost.\n\n**posteriordb/InferenceGym:** This is a great suggestion. Another reviewer also asked for more experiments with NUTS. Given the limited time, we have now implemented NUTS on several other datasets from our paper, and two datasets from posteriordb (Pilots and Soil-Incubation), which were in the list of challenging problems for NUTS in Table 1 of [1]. For the posteriordb problems we re-defined the logpdfs in JAX, and used the data from the JSON files at the posteriordb repository. Below we present the improvements in ESS/sec when implementing NUTS with our FSM (vs. BlackJax NUTS) on these datasets:\n\n- (Experiment 7.2) Real Estate: 3.14x\n- (Experiment 7.4) Google Stock: 3.36x\n- (Experiment 7.4) Predator Prey: 1.55x\n- (posteriordb) Soil: 2.42x\n- (posteriordb) Pilots: 0.91x\n\nWe find substantial speed ups in ⅘ cases, and larger speed-ups than the FSM variant of EllipSlice/TESS across all datasets we implemented for both methods. We note that because most of these datasets have highly variable local geometries, computational resources did not average out much across chains after 1000 samples. We expect that when sampling even longer chains, these speed-ups would be even larger. The Pilots dataset has a relatively cheap log-likelihood and less variation in the number of integration steps used, and so the gains in avoiding synchronization costs are offset by the increased cost per block of the FSM.\n\n**Pedagogical Example:** We agree that walking through the effect of synchronization and using our FSM on a simple example would add further clarity to the paper. We had wanted to do this in the submitted version but could not due to space constraints. We propose to use part of the additional page available to include this in Section 2 for the Elliptical Slice Sampler, and return to it again in Section 3 after presenting the FSM.\n\n**Switch Function Definition:** Thank you for the suggestion, we will include an explanation in the final version.\n\n[1] Magnusson et al. (2024).  “posteriordb: Testing, Benchmarking and Developing Bayesian Inference Algorithms”, arXiv preprint arXiv:2407.04967.", "claims_and_evidence": "This work makes careful and well-supported claims regarding the expected improvements of the proposed implementation strategy. The authors substantiate their arguments with rigorous yet concise theoretical analysis and thorough empirical validation.\n\nI find their reasoning compelling and fully agree with their conclusions.", "methods_and_evaluation_criteria": "yes, they make sense. \n\nI personally think a better metric to compare the efficiency gain is the number of synchronization steps per sec or per 100 sample.  Though, this might not be straightforward to log in the run time. \n\nAlso, I'm bit surprised that one FSM relative efficiency gain is only 3/4 times faster than the naive vectorization (as shown in Fig 6); I'm expecting a much more significant difference. See the question section for details.", "theoretical_claims": "I checked the correctness of the proof. The theoretical claims operate under very reasonable assumptions, and the proof is sort eluded in the maintext.", "experimental_designs_or_analyses": "The selection of metrics, choice of competitors, and Bayesian examples all seem well thought out and appropriate.\n\nHowever, for a fundamental idea like this (which I hold in very high regard), I believe a more comprehensive empirical evaluation would further strengthen the work. I recommend conducting additional evaluations on a broader set of models from PosteriorDB or TensorFlow Probability’s Inference Gym to provide more extensive validation.", "relation_to_broader_scientific_literature": "This work provides a thorough literature review of existing efforts to enable MCMC on modern hardware. To my knowledge, it takes a completely different approach—one that should absolutely be incorporated into modern probabilistic programming languages (PPLs). Instead of modifying the algorithm itself, this work focuses on optimizing its implementation, making it the first to do so in this manner.", "essential_references_not_discussed": "Not to my awareness.", "other_strengths_and_weaknesses": "n/a", "other_comments_or_suggestions": "- This could be due to my personal taste or lack of expertise in this field. However, I would appreciate a pedagogical example (e.g., ellipitical slice sampler with two chains) showcasing step by step the difference in the synchronization betweem the naive vectorization and the FSM variant.   To buy more space, I personally think it's okay to move section 5 to apdx.\n\n- Readers with less familarity with jax could benefit from a proper definition of the `switch` function in Algorithm 2.", "questions_for_authors": "- I'm bit surprised that one FSM relative efficiency gain is only 3/4 times faster than the naive vectorization (as shown in Fig 6) given that the drastic difference in iterations per sample.   Why doesn't walltime and ESS/sec doesn't scale in the same rate? (My guess is that for the walltime, it's dictated by the slowest chain, while iter/sample is controlled by the average speed across the chains.) \n\n- In the begining of section 4, the authors mentioned about the limitation of FSM vectorization that all branches would have to evaluated.  If my understanding is correct, this would be the same for the naive vectorization as well? As long as one call `vmap` or `@jit`, all branches of control flows will be evaluated? So in my perspective, regardless of the number of steps needed to obtain n samples, the performance of FSM vectorization will be lower bounded by the naive vectorization?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "uaybbFuOVv", "reviewer": "Reviewer_3kkK", "summary": "With the advancement of AI infrastructures, it is increasingly interesting to scale algorithms up with parallelism. Nevertheless, it is not efficient to run parallel MCMC with naive automatic vectorization (e.g., vmap), due to a varying execution time at each sampling step across different chains. The current approach will wait for the slowest chain at each step, while this paper proposes to remove this synchronization problem (usually caused by while-loops) by rewriting an MCMC algorithm as transitions in a finite state machine (FSM). With this modification, we can run parallel MCMC as parallel transition functions in the FSM, where different chains could be in different states, and synchronization will only happen once in the end. On the other hand, the proposed FSM based parallel MCMC also introduces additional control flow cost because codes for all states have to be executed every step no matter what state a chain is in. \n\nBased on the idea, this paper develops a theoretical framework for analysing such parallel MCMC algorithms. Namely, the expected running time of FSM v.s. standard parallel MCMCs will depend on the number of chains, the number of samples per chain, the convergence of the execution time per step, the \"organization\" of the FSM, etc. The takeaway is that the running time of standard parallel MCMC will converge to the slowest chain while the running time of FSM parallel MCMC will converge to a single chain scaled by the additional cost of control flow. \n\nTo reduce the cost of control flow, the paper proposes two techniques, called step bundling and cost amortization. Step bundling allows transitions to be combined in some cases. Cost amortization makes sure that an expensive function (usually the log density function) will only be executed once per FSM step. \n\nThe experiments test different aspects of the FSM parallel MCMC. FSM is the most helpful when the distribution of one-step sampling time is skewed. When implemented, FSM easily outperforms standard parallel MCMC with broad implementations of Delayed Rejection MH, elliptical slice sampling, HMC-NUTS and transport elliptical slice sampling, when the # of chains is sufficient. In addition, the control flow optimization is helpful on top of FSM and may introduce substantial speedups.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your diligent review - we are glad you believe our method is a worthwhile addition to modern MCMC libraries. Below we address your questions.\n\n**Theoretical framework:** Eq(3) and eq(4) hold without taking expectations under appropriate theoretical conditions, but we appreciate the symbol is doing a lot of heavy lifting in our effort to convey the basic ideas. With more standard probability theoretic notation, the exact form of eq(3) is $$\\sum_{i=1}^n\\max_{j \\leq m}N_{i,j}=n\\mathbb E\\max_{j \\leq m}N_{\\infty,j} + \\mathcal O_p(\\sqrt{n})$$ which holds under an appropriate Central Limit Theorem or concentration inequality on the sequence $(\\max_{j \\leq m}N_{i,j} : i \\geq 1)$. Since $0 < n \\mathbb E\\max_{j \\leq m}N_{\\infty ,j}  = \\mathcal O(n)$ in our setting, as $n$ grows the stochastic error term becomes negligible. \n\nFor eq(4), the exact version under equivalent conditions is \n$$\\max_{j \\leq m}\\sum_{i=1}^nN_{i,j}=  \\max_{j \\leq m} n \\mathbb EN_{\\infty,j} + \\mathcal O_p(\\sqrt{n})$$ \nwhich holds using the same arguments as eq(3) for the sum inside the max above, and using the fact that the max over finitely many $\\mathcal O_p(\\sqrt{n})$ variables is $\\mathcal O_p(\\sqrt{n})$.\n\nThis is all made rigorous in Theorem 4.1, where we use an approach based on Hoeffding bounds for Markov chains. \n\nWe appreciate the reviewer bringing this to our attention. We are happy to replace eq(3) and eq(4) by their reformulations above and modify the text accordingly.\n\n**Section 4 explanation:** We agree that more details on how the quantities in Section 4 hold would improve clarity.\n\nThe states $S_1,...,S_K$ exist in Algorithm 1 inside `sample` (recall that our FSM decomposes `sample` into contiguous states/blocks $S_1,...,S_K$ and defines `step` as a single call of one of these states). A single call of `vmap(sample)` for $m$ chains executes all $S_l$ (for $l \\neq k$) once, and $S_k$ (the iterative state) $\\max_{j \\leq m}N_{i,j}$ times, where $N_{i,j} =$ \\# while loop iterations to get  sample $i$ for chain $j$ and the max reflects that the while loop executes until it terminates for all chains.\nAssume the cost of state $S_j$, when executed for a batch of $m$ chains using `vmap`, is $c_j(m)$. In this case, the cost of calling `vmap(sample)` to get sample $i$ for all chains is $$\\sum_{l \\neq k}c_l(m) + c_k(m)\\max_{j \\leq m}N_{i,j}$$ Averaging this cost over $n$ samples gives eq(6) with $A_0(m) = c_{\\neg k}(m)=\\sum_{l \\neq k}c_l(m)$ (we appreciate $c_{\\neg k}$ should have been made clear) and $B_0(m) = c_k(m)$.\n\nFor eq(7), we follow similar steps with the main difference that now `vmap(step)` executes every block for all chains. This means the cost of a single FSM step is $\\sum_{j =1}^Kc_j(m)$. The overall cost of sample $i$ for chain $j$ is therefore $A_0(m) + B_0(m) N_{i,j}$, where $A_0(m) = (K-1)(\\sum_{j =1}^Kc_j(m))$ and $B_0(m) = \\sum_{j =1}^Kc_j(m)$. In the paper we multiply by $\\alpha$ to later motivate amortization. Averaging over $n$ samples gives eq(7).\n\n**Additional NUTS Experiments** Since another reviewer asked us to implement on datasets in the `posteriordb` database, in the limited time we chose to implement NUTS on (i) the Real-Estate dataset (Experiment 7.2), (ii) two datasets from Experiment 7.4, and (iii) two known challenging `posteriordb` datasets (`pilots` and `soil`). We used 128 chains + 1000 samples per chain + 400 steps pre-tuning for step size + mass matrix via the window adaptation algorithm. To save space, below we report the improvement in ESS/Sec using the FSM:  \n\n- Real Estate: 3.14x\n- Google Stock: 3.36x\n- Predator Prey: 1.55x\n- Soil: 2.42x\n- Pilots: 0.91x\n\nWe find substantial speed ups in ⅘ cases, and larger speed-ups than the FSM variant of Ellip-slice/TESS across all datasets we implemented for both methods. We note that because these datasets have highly variable local geometries (except real-estate GPR), computational resources did not average out across chains after 1000 samples. We expect that when sampling even longer chains, these speed-ups would be even larger. The Pilots dataset has a relatively cheap log-likelihood and less variation in the number of integration steps used, and so the gains in avoiding synchronization costs are offset by the increased cost per block of the FSM.\n\n**Using extra samples**: We also had the idea to let chains that finish early collect more samples while waiting for other chains. The problem is that this biases the samples, as chains initialized in regions that are easier to sample from are over-represented. Importance weighting could resolve this for estimating expectations, but we leave this to future work. Enforcing n samples per chain still allows us to avoid the inefficiency from chains waiting at every iteration with `vmap`.\n\n**Typos**: Thank you for pointing those out, we will amend them in revisions. In line 8 of Alg 5 `g_flag` is not meant to be returned - the flag is an output of the state functions, and is only used in `amortized_step`.", "claims_and_evidence": "The empirical evidences are strong but the theoretical evidences have several presentation issues. However, I do believe all claims are correct. See below for details.", "methods_and_evaluation_criteria": "The proposed method solves a key problem in efficient MCMC sampling. The problem of synchronization is part of why MCMC were mainly run with few chains on small machines. The solution of this work is neat and makes sense from the perspective of parallel computing systems. I would love to see the solution to be implemented in every modern MCMC libraries. \n\nThere are two axes of the benchmark: the MCMC algorithms and the models. The work has a broad selection of benchmark MCMC algorithms (delayed rejection MH, elliptical slice sampling, HMC-NUTS and transport elliptical slice sampling), which are all impactful in practice. The benchmark models include both simulated and practical models, covering difficult cases such as high correlations and funnels. I especially like the arguments of acceptance rate of at least one chain in parallel MCMC (line 436), which highlights the necessity of FSM based approaches.", "theoretical_claims": "My major issue with this paper is about the theoretical framework. I believe they are correct, but there are also some obstacles that keep it from being checked thoroughly. \n\n- Equations (3) and (4) work with expectations. I do not see how (4) is true given the information up until this part. More speficifically, my expression for (4) is $\\mathbb{E}[\\max_{j\\le m}\\sum_{i=1}^nN_{\\infty,j}^{(i)}]$. The gap might be justified with Theorem 4.1, but some connections need to be made.\n\n- The work refers to Algorithm 1 as \"standard design\" and Algorithm 3 as \"FSM design\". Both algorithms are analyzed based on the executing time of each state function in the FSM design. There is a disconnection between the state functions and Algorithm 1 as there are no states in Algorithm 1. To make sense, I am expecting an explanation of how much a step in Algorithm 1 costs as a function of $c_1, c_2,..., c_K$.\n\n- Following the above point, I find the notations around line 220 confusing. $c_{\\neg k}$ is not defined, and even if it is, I am not expecting $A_0$ to depend on a specific $k$. The presentations of $A_0, B_0, A_F, B_F$ are all not justified.", "experimental_designs_or_analyses": "I checked the settings of all experiments and they all look reasonable to me.\n\n-  Among all the MCMC algorithms, I am using NUTS the most. The NUTS experiment in 7.3 assumes identity mass matrix which looks different from standard approaches. But it is chosen to demonstrate how FSM helps the synchronization problem. However, I believe there should be more experiments on NUTS for a paper like this. It would be much more interesting to try NUTS on the models of Table 1.", "relation_to_broader_scientific_literature": "As the paper also points out, the synchronization problem is well documented in the literature (BlackJax,2019;Sountsov et al.,2024;Radul et al.,2020). The paper borrows classic ideas of finite state machines from computer systems (Hopcroft et al.,2001), and shapes it with problems of MCMC sampling.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "I believe there should be another change of mind about parallel MCMC following this work. The synchronization problem caused some chains to run \"in vain\" waiting for the slowest chain. With the FSM design, however, the waiting by the faster chain IS useful for sampling. For example, if we set $n=1,000$, the slowest chain will generate $1,000$ samples while the fastest chain might have generated $1,100$ samples. They shall all be mixed together for estimation. In practical settings, we can set a limit of # of transition cycles instead of # of samples to make more sense.", "other_comments_or_suggestions": "A list of typos I can find:\n\n- Line 121 right: $k\\in\\\\{0,1\\\\}$ -> $k\\in\\\\{1,2\\\\}$\n- Line 159 right: $k$ starts from 1 so probably $(1,z_1)$\n- Line 212 right: $(n,m)$ and $(m,n)$ are used inconsistently\n- Line 8 Algorithm 5: should $g\\\\_flag$ be returned?", "questions_for_authors": "My points can be condensed into the questions:\n\n- Is there a more intuitive justification of (4), as well as (6) and (7)?\n- How good is the performance of NUTS-FSM compared to NUTS on the models in Table 1?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a novel and efficient parallelization scheme for MCMC algorithms. A known limitation of the ${\tt vmap}$ function is its inefficiency when the function to be parallelized includes one or more while loops, as execution is delayed until all parallel processes have completed. The authors address this issue by leveraging finite state machines (FSMs). They demonstrate how common MCMC algorithms involving while loops can be represented using FSMs, and show that combining this representation with ${\tt vmap}$ enables more efficient parallel execution. Their claims are supported by both a theoretical complexity analysis and empirical results.

Most reviewers agree that this is an interesting and valuable contribution to the community. I share this view and recommend acceptance.

## Remark on Theorem 4.1
In the proof of Theorem 4.1, the authors assume that the Markov chain under consideration admits an absolute spectral gap. This assumption should be explicitly stated in the theorem, and a clear definition of the absolute spectral gap should be included in the supplementary material.
Decision: Accept (spotlight poster)