Parallelizing MCMC Across the Sequence Length

We thank the reviewer for their consideration of the paper and helpful feedback. In the following, we aim to address the reviewer’s questions about DEER, theoretical claims, and clarity of original contributions. Furthermore, we comment on the noted weaknesses, and report results applying the approach to larger-scale problems and adaptive MCMC samplers.

Elaboration on DEER algorithm (Question 1)

We agree with the reviewer that including more details about DEER to make the background self-contained will help with clarity of the paper. We have expanded the background to incorporate this feedback, notably by adding a pseudocode algorithm block that demonstrates how the DEER family of algorithms iteratively update the entire sequence of MCMC states until numerical convergence.

Theoretical Claims and Guarantees (Question 2, Weakness 3)

Our claims are typically supported by propositions from Lim et al., 2024 and Gonzalez et al., 2024. In particular, Proposition 1 from Gonzalez et al., 2024 shows that global convergence is achievable even when approximate Jacobians are substituted for the true Jacobians. Therefore, even in stochastic quasi-DEER where we have substituted in a stochastic estimate of the Jacobian, global convergence remains achievable. The sublinear scaling claims are based on the theoretical complexity of the work-efficient parallel scan ($\mathcal{O}(\log T)$ for length $T$), such that a relatively small number of parallel iterations is faster than sequential evaluation. We have clarified these statements in the paper by stating and citing the relevant propositions and assumptions made in each proposition.

For claims about quadratic convergence, we rely on showing the equivalence to Newton’s method, the proof of quadratic convergence in Appendix A.3 from Lim et al., 2024, and the general theory of Newton’s method. We will add this discussion to the paper. Importantly, we note that a complete theory of the number of iterations required for DEER to converge is not available, as we do not know when the quadratic convergence criteria hold, among other factors. Deriving such a theory and connecting it to the MCMC applications shown here is an important direction of future work, but it is outside the scope of this paper. However, as the reviewer notes the empirical results provide strong evidence for convergence in tens of iterations in the considered settings.

Clarity of contributions and Novelty (Question 3, Weakness 4)

We appreciate this feedback and will aim to more clearly state the novel contributions of this paper and how it is different from previous work. We will emphasize that the bulk of work on parallelizing diffusion model sampling relies on Picard iterations, which we believe works in that setting with small step sizes but can struggle to efficiently converge in MCMC applications with larger step sizes. In fact, applying the methods proposed in this paper in diffusion model examples is a promising future direction. Next, we will elaborate on various alternative ways parallel computation has been used for MCMC, though to our knowledge a general purpose procedure for parallel evaluation of standard MCMC methods across the chain length is novel.

Scalability (Weakness 2)

The reviewer states that scalability to higher dimensions or long chains remains a concern. We agree that these are important considerations. In the following, we report an application to a higher-dimensional problem and discuss additional factors for scaling.

First, we tested our approach in a sentiment classification model that predicts a binary positive or negative sentiment given a 768-dimensional LLM embedding of a text sequence. We used 1024 examples from the IMDB review dataset for this problem. We used parallel or sequential MALA to sample from the 768D posterior of this Bayesian logistic regression model. We found that parallel MALA simulated 4 chains of 4096 samples from this model about 3.66x faster than sequential MALA (mean sequential time = 0.448s, mean parallel time = 0.122s). This importantly demonstrates feasibility of the approach in higher dimensions. We will include this result in the final version of the paper.

Next, we address a few points regarding scaling. First, in practice we found that the stochastic quasi-DEER method rarely requires more iterations than the original diagonal quasi-DEER method, and the use of stochastic estimators is typically not a problem for scaling. Second, after submission we have also adopted a sliding window approach (akin to Shih et al., 2023) for scaling to longer sequences. The window can be set to the maximum sequence length supported efficiently in memory for the given batch size.

Significance of Advance and Adaptive MCMC (Weakness 4)

We focused on non-adaptive methods for this paper, as Gibbs, MALA and well-tuned HMC are useful in their own right for many inference and machine learning applications. We believe there are potential immediate applications of this work for machine learning systems that leverage Langevin dynamics, for approximate inference (i.e. via the proposed early stopping procedure), and for training algorithms that rely on MCMC sampling such as Monte Carlo EM. Furthermore, we envision that this work may open up new opportunities to apply MCMC sampling in settings where it was previously too time-consuming to run a sequential chain. Finally, the proposed stochastic quasi-DEER algorithm solves a general challenge for quasi-DEER, as computing the diagonal Jacobian can generally be time-consuming and memory intensive. The stochastic estimate can be computed in a single forward pass of the dynamics function that jointly evaluates the function at a specified input and returns the Jacobian-vector product (e.g. using jax.jvp). We therefore believe it will be generally useful for nonlinear sequence model training. For those reasons we believe this paper is more than an incremental advance over previous work.

However, we certainly agree with the reviewer that extending the approach to adaptive MCMC samplers is an important and exciting direction of future work. We tested the feasibility of parallelizing NUTS using Blackjax’s implementation of NUTS targeting the Bayesian logistic regression posterior for the German Credit dataset. On this model, we found that our proposed methods can effectively parallelize NUTS sampling under certain conditions and the resulting sampler was faster than sequential NUTS evaluation. With a batch size of 4 chains and 1024 samples per chain, the parallel sampler was about 3.98x faster than sequential sampling (mean sequential time = 1.139s, mean parallel time = 0.286s). These initial results, while preliminary, are a promising signal that the proposed approach can be extended to adaptive MCMC samplers such as NUTS. Furthermore, methods for dynamic HMC such as jittering the step size or number of leapfrog steps can be even more directly incorporated into our framework.

We thank the reviewer for their consideration of the paper and for their comments and questions. In the following, we clarify the MALA algorithm details, comment on the accuracy of the posteriors, present a new result demonstrating efficacy of the approach in a higher-dimensional problem, and address the reviewer’s question about multimodal posteriors.

MALA algorithm details

We greatly appreciate the question and feedback about our description of the MALA algorithm. There is a typo in line 94 of our writeup of the stop gradient trick that likely caused this confusion. The application of the stop gradient trick in line 94 should be written as $g \leftarrow \sigma(\tilde{g}) + \mathrm{stop gradient} (\mathbf{1}(\tilde{g}>0) - \sigma(\tilde{g}))$. Specifically, we introduce a gating variable $g$ that is binary in the forward evaluation of the function such that it implements a true, discrete accept-reject step. When computing the Jacobian, we interpolate between the current and proposed states using a sigmoid. Therefore, the ground truth sequential samples are true samples from a MALA sampler and the parallel sample sequence converges to the same sample sequence, up to numerical tolerance (in fact we tune the step size to ensure a reasonable rejection rate and observe rejected samples in practice). We have revised the manuscript with the corrected equation and improved explanation of this point.

Accuracy of algorithms

The reviewer states none of the experiments demonstrate the accuracy of the approximations. Here we hope to respectfully describe how we believe this is a misconception. Outside of the early-stopping results presented in Figure 6B, the parallel MCMC samplers converge to the ground-truth sequential samples, which are not approximate. Additionally, we put considerable effort into showing the quality of samples returned from the parallel MCMC samplers. We elaborate on both these points below.

First, we note that the sequential sampler in each example is not approximate, as it generates a set of ground truth samples obtained from running a Gibbs, MALA, or HMC sampler with manually tuned hyperparameters targeting a posterior of interest (we note the Gibbs sampler actually has no hyperparameters). Each example shown targets a commonly-used posterior for evaluation of MCMC algorithms. We showed for each example how samples from the parallel version of the algorithm converge to the ground truth sequential samples up to the numerical tolerance.

Second, for the MALA example we importantly used the maximum mean discrepancy (MMD) metric to evaluate the quality of the parallel and sequential samples as a function of wall-clock time. The MMD metric compared the simulated samples to a “gold-standard” set of 100K MCMC samples obtained using the No-U-Turn Sampler (NUTS). Both sequential and parallel MALA samplers achieved lower MMD values as more samples were obtained, which additionally demonstrated the accuracy of the samplers.

We will amend the manuscript to make these points more clear. Furthermore, we will also illustrate how the resulting posterior mean of the MALA model provides accurate predictions for the Bayesian logistic regression model.

Dimensionality of experiments and extending to higher dimensions

We agree about the importance of testing the approach on higher-dimensional distributions, as the German Credit example has 24 features. Here, we will describe an experiment demonstrating the approach in a significantly higher-dimensional problem targeting a posterior with 768 dimensions.

Specifically, we tested parallel MALA in a sentiment classification model that predicts a binary positive or negative sentiment given an LLM embedding of a text sequence. We used 1024 randomly selected examples from the IMDB Sentiment Analysis dataset for this problem. We computed a 768 dimensional embedding of each IMDB text sequence using Gemini Embedding and used a Bayesian logistic regression model to predict the sentiment from the embedding. Hence, the resulting model targets a higher-dimensional posterior with 768 dimensions.

We found that parallel MALA simulated 4 chains of 4096 samples from this model about 3.66x faster than sequential MALA (mean sequential time = 0.448s, mean parallel time = 0.122s). This importantly demonstrates feasibility of the approach in higher dimensions and we will include this result in the final version of the paper. Furthermore, while it is important to see successful scaling of the algorithm to higher dimensions, we also want to emphasize that it is not uncommon to have tens of dimensions for MCMC applications. Parallelizing MCMC for those problems is important in its own right.

Multimodal posteriors

The efficacy of this approach for multimodal posteriors is an important question, as outside of the Gibbs sampler we primarily focused on standard or hierarchical Bayesian GLMs in the submitted manuscript. First, since the parallel samplers converge to the sequential sampling states, the efficacy of exploring multimodal distributions depends on the ability of the sequential sampler to explore such distributions. However, we note that DEER is not limited to small numbers of samples. The parallelization actually can enable efficient generation of much longer MCMC sequences, which may help algorithms explore multimodal distributions.

Here we demonstrate applicability of the samplers to multimodal distributions in a toy example. We tuned MALA to target a mixture of 4 Gaussians in two dimensions such that it explored all modes. Using quasi-DEER, we found parallel MALA could evaluate a 10K length sequence with 41 parallel iterations, converging to the ground truth sequential trace. We have revised the manuscript to include this example in the appendix.

Conclusion

We hope to have addressed the reviewer’s concerns regarding the MALA algorithm, accuracy of the samplers, and scalability of the algorithm to higher-dimensions, and the reviewer’s questions regarding multimodal distributions. If the reviewer has any remaining concerns or questions, we are happy to address them in the discussion phase.

We thank the reviewer for their consideration and feedback. We discuss below how we have improved the background section to provide more details on DEER. Additionally, we report two new experiments we performed that test the approach for 1) adaptive MCMC samplers and 2) higher-dimensional problems.

Additional Background on DEER

We appreciate the suggestion to add more background on DEER to the methods, as we agree that this will improve clarity. We have revised the paper accordingly. In particular, we added a pseudocode algorithm block showing the high-level steps of this approach. This more clearly shows how the algorithm initializes a guess for the full sequence of states and iteratively updates the entire sequence using a variant of DEER, which is parallel-in-time. We have also added additional explanatory details to the text, including considerations for computing the Jacobians.

Extending approach to adaptive MCMC Samplers

An important direction of future work is extending the approach to adaptive MCMC samplers such as the No-U-Turn Sampler (NUTS). We tested the feasibility of parallelizing NUTS using Blackjax’s implementation of NUTS targeting the Bayesian logistic regression posterior for the German Credit dataset. On this model, we found that our proposed methods can effectively parallelize NUTS sampling under certain conditions and the resulting sampler was faster than sequential NUTS evaluation. With a batch size of 4 chains and 1024 samples per chain, the parallel sampler was about 3.98x faster than sequential sampling (mean sequential time = 1.139s, mean parallel time = 0.286s). These initial results, while preliminary, are a promising signal that the proposed approach can be adapted to adaptive MCMC samplers such as NUTS.

Demonstration of approach in higher-dimensions

Finally, we also want to mention an additional experiment we performed investigating the method in a higher-dimensional problem. We tested our approach in a sentiment classification model that predicts a binary positive or negative sentiment given a 768-dimensional LLM embedding of a text sequence. We used 1024 examples from the IMDB review dataset for this problem. We used parallel or sequential MALA to sample from the 768D posterior of this Bayesian logistic regression model. We found that parallel MALA simulated 4 chains of 4096 samples from this model about 3.66x faster than sequential MALA (mean sequential time = 0.448s, mean parallel time = 0.122s). This importantly demonstrates feasibility of the approach in higher dimensions. We will include this result in the final version of the paper.

We thank the reviewer for their consideration and comments. In the following, we discuss the feasibility of extending the framework to adaptive MCMC samplers such as the No-U-Turn Sampler (NUTS). We also comment on the relationship between the wall-clock speedup and hardware and report results on applying the framework to a higher-dimensional problem.

Adaptive MCMC Samplers (Question 1, Weakness 1)

Extending the proposed framework to adaptive MCMC samplers is an important and exciting direction of future work. Here, we discuss considerations for variations of adaptive MCMC and also report the results of an initial test of applying the approach to parallelizing NUTS targeting the Bayesian logistic regression posterior on the German credit dataset. We will include the following discussion of adaptive MCMC methods in the final version of the paper.

NUTS introduces additional control flow for building the set of candidate points. As the proposed algorithms in this paper leverage the Jacobian of the update step to propagate information, this potentially complicates extending the approach to NUTS. Nonetheless, we aimed to quickly test the feasibility of this approach. We used the Blackjax package’s NUTS implementation [Cabezas et al., 2024], which is compatible with a forward-mode Jacobian-vector product out of the box because it uses jax.lax.cond. This propagates gradient information along the evaluated branch of the algorithm but notably does not differentiate through the predicate that determines which branch to evaluate. We did not modify the code to attempt to improve the Jacobian signal (for example, by substituting in differential relaxations for Jacobian computation).

On the Bayesian logistic regression model for the German credit dataset, we found that the approach can effectively parallelize NUTS sampling under certain conditions and the resulting sampler was faster than sequential NUTS evaluation. On a GPU with a batch size of 4 chains and 1024 samples per chain, the parallel sampler was approximately 3.98x faster than sequential sampling (mean sequential time = 1.139s, mean parallel time = 0.286s). These initial results, while preliminary, are a promising signal that the proposed approach can be adapted to NUTS.

Aside from NUTS, a simpler approach to avoid potential HMC pathologies with a fixed step size and trajectory length is to randomly jitter either the step sizes or number of leapfrog steps across sample iterations (sometimes called Dynamic HMC). While we have not tested this, these methods are straightforward to incorporate into our approach. In particular, the jittered step sizes or number of leapfrog steps can be presampled and wrapped up into the time-varying function $f_t$. Notably, ChEES-HMC [Hoffman et al., 2021] is an adaptive HMC algorithm that uses dynamic trajectory lengths after tuning parameters during a warm-up phase, and our approach could be applied using the tuned parameters after an initial warm-up.

Relationship between wall-clock speedup and hardware (Weakness 2)

Here we discuss the theoretical efficiency of the algorithm and how it scales with additional hardware resources, considering the reviewer’s comment on wall-clock speedup scaling with hardware. We focus on the work-efficient parallel scan used to evaluate the linear recursion. If $P$ is the number of available parallel processors, then the time complexity for evaluating a length $P$ sequence is $\mathcal{O}(\log P)$. If the number of parallel processors is doubled to $2P$, then we can then evaluate a length $2P$ sequence in $\mathcal{O}(\log 2P)$ time. This in fact implies a beneficial theoretical scaling of the algorithm with additional hardware for sequence length, where for $P>2$ doubling the processors enables evaluating twice the length of a sequence in less than twice the amount of time.

Demonstration of approach in higher-dimensions

Finally, we report an additional experiment we performed investigating the method in a higher-dimensional problem. We tested our approach in a sentiment classification model that predicts a binary positive or negative sentiment given a 768-dimensional LLM embedding of a text sequence. We used 1024 examples from the IMDB review dataset for this problem. We used parallel or sequential MALA to sample from the 768D posterior of this Bayesian logistic regression model. We found that parallel MALA simulated 4 chains of 4096 samples from this model about 3.66x faster than sequential MALA (mean sequential time = 0.448s, mean parallel time = 0.122s). This importantly demonstrates feasibility of the approach in higher dimensions. We will include this result in the final version of the paper.