Efficiently Vectorized MCMC on Modern Accelerators

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.

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.

For eq(4), the exact version under equivalent conditions is $\max_{j \leq m}\sum_{i=1}^nN_{i,j}= \max_{j \leq m} n \mathbb EN_{\infty,j} + \mathcal O_p(\sqrt{n})$ which 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})$.

This is all made rigorous in Theorem 4.1, where we use an approach based on Hoeffding bounds for Markov chains.

We 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.

Section 4 explanation: We agree that more details on how the quantities in Section 4 hold would improve clarity.

The 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. Assume 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)$.

For 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).

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:

• Real Estate: 3.14x
• Google Stock: 3.36x
• Predator Prey: 1.55x
• Soil: 2.42x
• Pilots: 0.91x

We 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.

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.

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.

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.

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.

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.

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

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:

• (Experiment 7.2) Real Estate: 3.14x
• (Experiment 7.4) Google Stock: 3.36x
• (Experiment 7.4) Predator Prey: 1.55x
• (posteriordb) Soil: 2.42x
• (posteriordb) Pilots: 0.91x

We 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.

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.

Switch Function Definition: Thank you for the suggestion, we will include an explanation in the final version.

[1] Magnusson et al. (2024). “posteriordb: Testing, Benchmarking and Developing Bayesian Inference Algorithms”, arXiv preprint arXiv:2407.04967.

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.

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.

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.

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.

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:

"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".

$^*$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.

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:

"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)".

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 vmapdoes 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.

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:

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.

In more detail:

• 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).
• 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.

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)$.

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.

Other typos+ formatting suggestions: Thank you for pointing out the typos, formatting and improvements, we will amend those accordingly in the final version.