Parallelizing MCMC Across the Sequence Length
Authors: David M. Zoltowski
Skyler Wu
Xavier Gonzalez
Leo Kozachkov
Scott Linderman
Venue: NeurIPS 2025 poster ()
Keywords: probabilistic methods
Markov chain Monte Carlo
parallel algorithms
Primary area: probabilistic_methods
TL;DR: We parallelize MCMC over the sequence, yielding more than order of magnitude wall-clock speedup on sampling.

Abstract: Markov chain Monte Carlo (MCMC) methods are foundational algorithms for Bayesian inference and probabilistic modeling. However, most MCMC algorithms are inherently sequential and their time complexity scales linearly with the sequence length. Previous work on adapting MCMC to modern hardware has therefore focused on running many independent chains in parallel. Here, we take an alternative approach: we propose algorithms to evaluate MCMC samplers in parallel across the chain length. To do this, we build on recent methods for parallel evaluation of nonlinear recursions that formulate the state sequence as a solution to a fixed-point problem and solve for the fixed-point using a parallel form of Newton's method. We show how this approach can be used to parallelize Gibbs, Metropolis-adjusted Langevin, and Hamiltonian Monte Carlo sampling across the sequence length. In several examples, we demonstrate the simulation of up to hundreds of thousands of MCMC samples with only tens of parallel Newton iterations. Additionally, we develop two new parallel quasi-Newton methods to evaluate nonlinear recursions with lower memory costs and reduced runtime. We find that the proposed parallel algorithms accelerate MCMC sampling across multiple examples, in some cases by more than an order of magnitude compared to sequential evaluation.

Parallelizing MCMC Across the Sequence Length
David M. Zoltowski" Stanford University tzoltow@sttanfrd edu
Skyler Wu" Stanford University s akyleer@stanford.eed
Xavier Gonzalez
Stanford University
xavier180stanforct edu
Leo Kozachkov Brown University 1 1eokoz8@boow..c edu
Scott W. Linderman
Stanford University
scott www..eeeerrnnnttafrrr...d..
Abstract
Markov chain Monte Carlo (MCMC) methods are foundational algorithms for Bayesian inference and probabilistic modeling. However, most MCMC algorithms are inherently sequential and their time complexity scales linearly with the sequence length. Previous work on adapting MCMC to modern hardware has therefore fo- cused running many independent chains in parallel. Here, take alternative approach: we propose algorithms to evaluate MCMC samplers in parallel across the chain length. To do this, we build recent methods for parallel evaluation of nonlinear recursions that formulate the state sequence as a solution to fixed-point problem and solve for the fixed-point using parallel form of Newton's 's method. We show how this approach can be used to parallelize Gibbs, Metropolis-adjusted Langevin, and Hamiltonian Monte Carlo sampling across the sequence length. In several examples, we demonstrate the simulation of up to hundreds of thousands of MCMC samples with only tens of parallel Newton iterations. Additionally, we develop two new parallel quasi-Newton methods to evaluate nonlinear recursions with lower memory costs and reduced runtime. We find that the proposed parallel algorithms accelerate MCMC sampling across multiple examples, in some cases by more than an order of magnitude compared to sequential evaluation.
Introduction
Markov chain Monte Carlo (MCMC) algorithms sequentially generate samples from a target distribu- tion, with computational cost linear in the length of the chain [1-3]. While modern hardware like GPUs and TPUs can dramatically accelerate parallelizable algorithms, leveraging these advances for MCMC remains challenging. In particular, while parallel resources can be used to simulate many independent chains in parallel [4, 5J, the runtime remains linear the sequence length.
In this paper, we present a promising approach for parallelizing MCMC algorithms across the sequence length obtain sublinear chain-length complexity. The core idea is to adapt methods for the parallel evaluation of nonlinear recursions [6-8] to MCMC algorithms. These methods formulate the state sequence of nonlinear sequence model as the solution of an optimization problem and iteratively solve for the state sequence via parallel-in-time formulation of Newton's 's method.
We show that these methods can parallelize widely used classes of MCMC samplers across the chain length, including Gibbs sampling [9], the Metropolis-adjusted Langevin algorithm (MALA) [10], and Hamiltonian Monte Carlo (HMC) [11-13]. Across these examples, often find that dozens of parallel evaluations are sufficient to generate anywhere from thousands to hundreds of thousands of samples. We also propose two new variations of the parallel Newton method that reduce computational costs and improve efficiency for certain MCMC applications. With these approaches, we find that the proposed parallel algorithms can dramatically accelerate MCMC sampling.
"Equal contribution.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

A
Sequential Sampling
Parallel Sampling
B
Parallel Iter. 1
Parallel Iter. 10
40
iter.
20
X
Iter. 10
Parallel Iter. 25
Parallel Iter.
147
Iter. 50
500
1000
500
1000
-25
0
25
-25 25
Sample Index (t)
Sample Index (t)
X1
Sequential Samples Iteration
Parallel Samples Iteration
Figure 1: Parallel evaluation of 100K HMC samples targeting the Rosenbrock distribution. (A) Dark lines (black: sequential, blue: parallel) show simulated samples of the second coordinate for the first IK samples at three iterations each method. The light gray line the final sequential chain. (B) 100K samples after 10, 25, or 147 parallel iterations (blue), converging to the sequential trace (gray) by iteration 147.
The paper is organized as follows. We first review how Newton's method can parallelize nonlinear recursions, then show how this framework applies to Gibbs sampling, MALA, and HMC across the sequence length. For HMC, we additionally propose parallelizing only the leapfrog integration within each proposal. Next, we introduce scalability improvements to parallel Newton methods and novel block quasi-Newton variant adapted for parallelizing leapfrog integration. Finally, present experiments demonstrating wall-clock speedups of parallel MCMC sampling across multiple examples. In Figure 1, we illustrate our approach by simulating 100K HMC samples from high- curvature 2D distribution, where parallel HMC converges to the sequential trace in just 147 steps. Our implementation is available https: petppa//gittul com hwtp:/innndddnneeeeeeee aat.llallapaaaaaaaaa httppppaaallllemeeea
2 Background
Our work builds on and extends previous approaches for parallelizing linear and nonlinear recursions. Consider recursion, = fe(ss))), that when given an initial state So yields sequence of states, 81.7 (S1,..., ST) C TR. according to the transition functions, Jt TR. + R.D The obvious way to solve for the sequence of states is to evaluate the recursion sequentially. However, if fa is an affine function (i.c., fc(s = J( + u, for some matrix J, and vector uz), then the system can be evaluated in O(log I') time parallel machine using parallel (a.k.a. associative) scan [14]. The key insight is that a composition of updates, f++1 fc. still an affine function. With (TT) processors, one compute all pairs of updates in parallel, yielding sequence that is half long. By repeating this process (log 7') times, one can obtain the entire sequence in sublinear time. This parallel algorithm underlies many modern machine learning models for sequential data [15-17].
In contrast, nonlinear sequential processes such as recurrent neural networks and, for our purposes, common MCMC algorithms, are not directly parallelizable and instead are typically evaluated sequentially. However, Danieli et al. [6] and Lim al. [7] showed that we can view the states of nonlinear recursion as the solution of fixed-poins equation, and that a parallelized form of Newton's method can be used to solve for the state sequence. Lim et al. [7] and Danieli et al. [6] observed (4)
that iteratively linearizing the transition functions, ft. around guess for the state trajectory, $1.7- and evaluating the resulting linear system to obtain a new trajectory, $1: equivalent to using Newton's 's method on appropriately defined residual. More precisely, given Sp and fi.r, the residual the vector of temporal differences
r(si.) = vec([s) fi (so). .... ST (--(f(ST--))
and the Newton update is given by the Taylor expansion,
s (i+1) f(() J. (
(i+1)
s(21). where J. == 0ft (6)
(s/_1).
(2)

---

By rearranging terms it is clear this is linear dynamical system
Si (i+1)
(++1) J J + fe(s(_)) - J
with dynamics given by the time-varying Jacobians, J. and inputs, u. The Jacobians and inputs only depend on the states from the previous Newton iteration, and thus can be computed in parallel. Crucially, this lincar dynamical system can also be solved in (log I") time using the parallel scan algorithm described above. This parallel update is repeated until the state sequence converges within numerical tolerance 8. This algorithm was termed "DEER" by Lim et al. [7] and "DeepPCR" by Danieli et al. [6]. For clarity, we refer to it as "DEER" or "Newton's 's method."
parallel machine
The DEER algorithm requires computing, storing, and multiplying T Jacobian matrices J, of size D x D. Therefore, DEER must perform O(TD*) work and requires OTTD²) memory, which can be prohibitive in larger dimensionalities. To reduce computational complexity, Gonzalez et al. [8] proposed quasi-DEER, a quasi-Newton method that retains only the diagonal of the Jacobian matrices, d))g(J and evaluates the following linear recursion
(++1) diag(J s (i+1) fr(s(2)) diag (J:)
%
(4)
This reduces storage and matrix multiplication costs to O(TD). which generally improves wall-clock time and memory usage compared to DEER. However, computing the diagonal of the Jacobian via automatic differentiation still requires D passes through the function, which can be slow or memory-intensiie. Gonzalez et al. [8] thus propose hard-coding the diagonal of the Jacobian when possible. In Section 3.4, we propose an alternative that requires only single forward pass through the function and does not require care-coded diagonal Jacobians. In summary, DEER variants evaluate nonlinear sequence models by iteratively refining an initial state sequence using parallel-in-time updates. A high-level overview of this procedure is shown in Algorithm in Appendix A.
3 Parallelizing MCMC
In this work, we leverage parallel Newton methods to evaluate MCMC samplers. We set the sequence of functions fi:r to be the transitions of an MCMC sampler and So to be the initial state of the chain. We then iteratively solve for the resulting MCMC sampling sequence S1.7. In this section we describe the details of this approach for parallelizing reparameterized Gibbs sampling, MALA, and HMC, as well as several techniques for improving the efficiency and scalability of parallel MCMC.
3.1 Parallel reparameterized Gibbs sampling
Gibbs samplers partition the joint distribution of the target into conditional distributions that can be alternately sampled [9]. Here we focus on a subset of Gibbs samplers in which each conditional distribution is reparameterizable, such that it a deterministic and differentiable function of input randomness. Suppose we are targeting a distribution p(z) over random variable æ E R² Let 2; = (It. DI. I.." denote the state of the Gibbs sampler iteration We define f to be the Gibbs sweep that iteratively updates each coordinate given input noise, & E DR such that
at = ((x.-1, (((
(5)
Even though the MCMC transition kernel is fixed, the reparameterized update functions vary with time because the input noise, &., changes at each step. More specifically, let fa be the Gibbs update of coordinate conditioned on other coordinates.
Jt,d = fa (1 ... Lt,d-1, I1-1, It-1 Er,a),
(6)
which importantly does not depend on 21-1,d- Then ff is the composition of the D coordinate updates fa o ..... .... o fD. Observe that when put in the form of equation (5) this is a nonlinear
f == fi dynamical system with states, s; Dz and transition functions, fi, which vary in time due to the input randomness &,- Thus, we can directly apply the DEER algorithm to parallelize sampling.
3.2 aediepells-adjussed Langevin dynamics (MALA)
MALA generates proposals via Langevin dynamics and then accepts rejects the proposed state with a Metropolis correction [10]. For target distribution p(x) and step size e, the proposal is generated via
x x + e logp((x-1)) + 22EE.. & I(
3

---

and accepted with probability min(1, ( -x)pp( x(-1)} where q(x, x(-1) the proposal density of candidate x, given x;- Importantly, MALA can be cast nonlinear recursion amenable to parallelization by setting the state 8 x. The MALA update function f1( A MALA( (E.. ( depends on input random variables &, ~ ~N(0 I) and oe ~ (4(0,) 1). We define it as follows, with logistic function a and binary indicator function I(...):
function MALA(X,-1. (E, 2z3)
X, - x;-1 + cVx (xx 22c,,
a f min ((x))((xx-- (1 81)/p(x,-11 11/((( xx-1)))
log a log u
g f a (ğ) + stop_groeeeee grapggaaadnnnttn > 0) (ī))
> stop- moop-graieet trick accept-reject via gating variable
X, < gx + (1 g) X
return X
end function
Notably, the accept-reject step renders the MALA update non-differentiable. In our proposed MALA function, the exact MALA update is computed in the forward pass, while the Jacobian for DEER is obtained using the stop-gradient trick [see 18], which yields a differentiable relaxation of the accept-reject step. Critically, Proposition of Gonzalez et al. [8] guarantees that DEER and quasi-DEER still globally converge to the true sequential MALA trace in this setting, despite the approximate Jacobian and non-differentiable ft. That is, given shared input randomness, both parallel and classical sequential MALA return the same set of samples up to the numerical tolerance d.
3.3 Parallelizing Hamiltonian Monte Carlo Hamiltonian Monte Carlo (HMC) is prominent MCMC algorithm that uses gradient information with augmented momentum variables to efficiently explore the sampling space [11-13]. Each step of HMC consists of sampling new momenta, integrating Hamiltonian dynamics, and accepting or rejecting the proposed state. In the following, we present two approaches for parallelizing HMC that either: 1) parallelize across sampling steps; or 2) parallelize the leapfrog integration within each step. We focus on HMC with a fixed number of leapfrog steps and step size and discuss considerations for future work on parallelizing NUTS [19] and other adaptive HMC algorithms in Section 6.
Parallel HMC with Sequential Leapfrog Integration.
The most straightforward way parallelize HMC is to extend our MALA strategy and parallelize across HMC sampling steps. Here, the function St runs a sequential leapfrog integrator generate a new proposal. As before for MALA, we substitute a differentiable relaxation of the accept-reject step for the Jacobian. We used this method in Figure 1 and apply it logistic regression in Section 5.
Sequential HMC with Parallel Leapfrog Integration. Alternatively, could run HMC steps sequentially and parallelize the inner leapfrog integration loop. Let p(x) be the target distribution with DR Each HMC step samples momentum variable E TR'2 and then integrates Hamiltonian dynamics for I. steps using the leapfrog integrator with step size e. If the momentum is updated with half-step before integrating, the resulting integration step is
(8) for / 1= (1, L). After integration, the momentum is updated with negative half-step and then the candidate state is accepted rejected. To apply DEER to this problem, we set the state to the concatenation of the position and momentum states S. = |x,, vz] and define ft : RR2D - TR2D as the leapfrog integrator steps in equation (8). The complete algorithm is presented in Appendix A.2. This novel ability to parallelize leapfrog integration raises new considerations about the optimal step size and acceptance rates, which we discuss in Appendix B.3.
Xt-1 V Vi = V:-1 + Vx log)(x.)
The Jacobian of the leapfrog step has Block quasi-DEER for parallelizing leapfrog integration. the following block structure, where I is an identity matrix:
0f ((s--)) Os (st-1) = 2 logp(x;-1 + cv.-1) ID V2
Io
d
(9)
The diagonal quasi-DEEs algorithm discards information in the off-diagonal blocks. We therefore propose an alternative "block quasi-DEER" algorithm that employs the diagonal of each block of the Jacobian. For leapfrog integration, this approximate Jacobian is exact for the top row and takes into account interactions between the position and momentum states in the off-diagonal blocks. Crucially, the block quasi-DEER matrix form can also parallelized via an efficient linear recursion with (TDD) memory and (TD)) work. See Appendix A for additional details of this approach.

---

3,4 Scaling parallel MCMC
In this section, we propose several ways to make parallel MCMC more efficient by reducing the computational cost of each update or reducing the number of parallel Newton iterations.
A general stochastic approach for efficient quasi-DEER. Quasi-DEER [8] and our block variant require the diagonal of Jacobian, which is not always easily available in closed form. Computing the diagonal via automatic differentiation requires D passes through the function to compute the full Jacobian and retain only the diagonal elements, which is computationally and memory intensive for large systems. We therefore propose a stochastic quasi-DEER that is generally memory efficient and still enjoys global convergence guarantees due to Proposition 1 of Gonzalez et al. [8]. Our approach leverages a general stochastic estimator of the diagonal of the Jacobian [20, 21] based on Hutchinson's method [22] and given by the following, where each element z, Rademacher:
diag( J..) = Eg-RRad (z O (J(z)..
(10)
This expectation can be estimated via standard Monte Carlo using efficient lacobian-vector vector products.
The resulting stochastic quasi-DEER algorithm joins the family of memory-efficieee quasi-DEER algorithms proposed in Gonzalez et al. [8]. Empirically, we find that only one or a few Monte Carlo samples often suffice estimate diag(J J Notably, when using a single sample estimate the stochastic quasi-DEER algorithm requires only one forward pass through the function ft to evaluate it at the state S and compute the Jacobian-vector product J,z. This substantially improves over the D passes through the function required for automatic differentiation and matches the number of forward passes required by Picard iterations efficiency gains, and we expect it to be broadly useful for parallel evaluation of nonlinear recursions.
[23]. In our results, stochastic quasi-DEEE yields major
High quality approximate samples with early-stopping. The sample sequence generated via par- allel Newton iterations at convergence is equal to the sequence generated from sequential evaluation. up to numerical tolerance. However, we observe that intermediate sample sequences before conver- gence can still apparently produce high-quality samples. For example, consider the samples generated at the 25th parallel iteration in Figure 1. We therefore propose an approximate sampling scheme where samples are generated by early-stopping the parallel Newton iterations before convergence.
Orthogonal coordinate transformation for quasi-DEER. As quasi- relies on a diagonal approximation to the Jacobian, its performance can depend on the accuracy of this approximation. To mitigate this issue in some problems, we propose to reparameterize the system for quasi-DEER using an orthogonal coordinate transformation z, = Q s, with Q E Dx and Q Q = I. In this new set of coordinates, the dynamics of the system are given by
2, = Q (( - f((z)-1).
(11)
The resulting Jacobian in these coordinates is J, Q J.Q. Jf , is closer to diagonal than J the quasi-DEER approximation will be more accurate. At convergence, we map the transformed states back to the original coordinates via S. Qz,- These efforts to make Markovian systems more amenable to quasi-DEER are part of a broader theme in the community; for example, Farsang et al. [24] develops nonlinear RNN where the dynamics are constrained to be diagonal. The orthogonal change of coordinates is particularly effective for dynamical systems with real and symmetric (or near-symmetric) Jacobians. Importantly, this is the case for MALA. If we ignore the accept-reject step in MALA then the Jacobian is J, I + < log 1 Assuming continuity of second derivatives of log pxx,, this is a real and symmetric matrix that can be diagonalized with an orthogonal matrix Q.
Sliding window updates. To improve algorithm performance for higher- dinher-dimensionall problems, we adapted the sliding window technique proposed in Shih et al. [23]. In this approach, at each iteration we apply a DEER update within window of the sequence rather than across the entire sequence length. The window is initialized to start at the first time point. After each iteration, the window is shifted forward in the sequence to the first time point that has not yet converged within the tolerance.
4 Related Work
Parallelizing diffusions. A related line of work has developed parallelizable algorithms for sam- pling from diffusions. Shih et al. [23] used Picard iterations, which also cast the state sequence as a fixed point, to parallelize diffusion model sampling. Follow-up works include the Parareal algorithm [25] by Selvam et al. [26] and Anderson acceleration via triangular nonlinear equations by Tang al. [27]. In fact, recent work has shown that in this setting, Picard iterations can be interpreted
5

---

A
B
Iters Until Convg.
# Chains:
# Chains: 32
# Chains: 64
# Parallel Iters.
10
20
Time (sec)
1.5
100
1.0
10
50
0.5
0.0
111
10"
10
10
10a
105
10
10*
10
104
10
Samples
Samples
#
Samples
Samples
DEER (Original)
Memory-Efricent Quasi-DEER Sequential (CPU) Sequential (GPU)
Interpolation
Figure 2: Parallel Gibbs sampling. (A) The number of iterations for DEER and preconditioned memory- efficient quasi-DFER scales approximately logarithmically in the sequence length (solid line is median and shaded areas are 90% confidence intervals). (B) Quasi-DEFR is faster than sequential sampling on CPU or GPU across varying batch sizes, i.e. number of independent chains evaluated in parallel.
as a version of quasi- our work, Anari et al. [29] used Picard iterations to parallelize Langevin diffusion sampling. While these studies are highly relevant, there are important differences from our work. First, they focus on parallel simulation of SDEs, which differ from MCMC algorithms involving discrete accept-reject steps and coordinate-vv -wise updates. Next, Picard iterations achieve linear convergence, whereas our Newton-bard based approach attains quadratic convergence under certain conditions. Separately, Danieli et al. [6] further demonstrated the promise of parallel Newton methods for diffusion model simulation, showing theoretical speedups under the assumption of perfect Jacobian parallelization, though limited in practice by memory constraints. The stochastic quasi-DEER approach may alleviate such memory issues and aid parallelizing diffusion model sampling.
where the Jacobian replaced by the identity [28]. More closely related
Related work for parallel MCMC. Many works have sought to accelerate MCMC via parallel computing [30, 4]. The most related Grazzi and Zanella [31], which parallelized random-wan Metropolis sampling using Picard iterations. In contrast, we focus on Gibbs sampling and gradient- based MCMC algorithms and employ first-order Newton methods for parallelization. Beyond this, the most common approach is run multiple independent chains in parallel, which is supported by modern MCMC libraries [32-35] and theoretical analysis [36, 5]. However, some samplers like the No-U-Turn Sampler (NUTS), adaptive HMC algorithm [19] with additional control-fior -looo computations that complicate GPU acceleration, may not be as easily parallelizable across chains.
This motivated GPU-friendly adaptive HMC variants [37, 4] and finite-state machine formulations of MCMC algorithms that alleviate GPU synchronization issues [38]. Finally, many parallel-chain MCMC algorithms also share information across chains [39, 40].
Parallel computation in MCMC extends beyond batching independent chains. For Gibbs sampling, extensive work has focused parallelizing Gibbs update sequences within sample iteration [41- 43]. In puoti-proopoall MCMC, multiple candidate states are generated in parallel to enhance target exploration [44-49]. Pre-fetching approaches parallelize tasks such as proposal generation alongside other computations [50, 51, 30). Parallel-in-time algorithms have been developed for Bayesian smoothing in state space models [52-54] and are related to work. Finally, many approaches partition data into subsets and run parallelizable inference processes for each data subset [55-57, 30].
Results
Our experiments evaluate the number of Newton iterations to convergence, wall-clock time relative to sequential sampling, and sample quality for parallel Gibbs, MALA, and HMC across multiple prob- lems. Our implementations were in JAX [58] with wall-clock times measured post-JIT compilation. Unless otherwise noted, all runs used a single H100 GPU on a SLURM cluster.
5.1 Parallel Gibbs sampling for a hierarchical Gaussian model
We first demonstrate parallelization of reparameterized Gibbs sampler for the eight schools prob- lem [3, 59]. hierarchical Gaussian model of test scores 2s,n for school s and student with 20 students per school (see Appendix B.2 for additional details). We sampled synthetic data from this model with specified means and standard deviations per school, and used parallelized Gibbs sampler to draw samples from the 18-dimensional posterior distribution.

---

A Wall-Clock Times B Iters. to Convg. C Wall-Clock Speedup Multipliers of Parallel over Seq. MALA
Time (sec)
Iterations
80
60
Speedup
L 2K
1K
= 32K
16K
20
= 4K
64K
40
L 8K
10
20
10a
10
10a
B-1
B-2
8-4 8-8 B-16
B=32
Seq. Length (L)
Seq. Length (L)
Batch Size (8)
Figure 3: Parallel MALA performance using efficient quasi-DEER (Q-DEER), (A) Wall-clock times to sample chains of length L for batch size B = 2 for parallel MALA ("Par.") ") vs. sequential MALA ("Seqq"). (B) For batch size B = 2, the distribution/number of Newton iterations needed for cach chain to numerically converge. The light-blue band represents 60% confidence interval over 20x random seeds. The gray dots show the number of iterations needed for numerical convergence a particular seed as example. The dashed line represents maximum permitted number of iterations at given setting of L. Parallel MALA often converges in only a few dozen iterations, even when generating tens of thousands of samples. (C) Wall-clock speedup multipliers of parallel over sequential MALA with 90% confidence intervals over 20x random seeds. The gray dotted line marks multiplier of L.0, ie.. equal performance. For many batch sizes, Q-DEER generate numerically eqqivvaleet samples more than an order of magnitude faster than sequential sampling.
We investigated the convergence and wall-clock speed of DEER and efficient quasi-DEER on an A 100 GPU using 3-sample stochastic Jacobian diagonal estimate and diagonal preconditioning. We compared wall-clock times to sequential sampling baselines on GPU and CPU across various numbers of chains and chain lengths. We found parallel Newton methods converged rapidly for this problem, requiring tens of DEER iterations or 100-150 quasi-DEER iterations for chain lengths up to IM samples (Figure 2A). Furthermore, for batches of 32 or 64 chains, Quasi-DEER yielded the fastest wall-clock sampling times across all settings (Figure 2B), achieving approximately 2x faster sampling than sequential methods.
5.2 Parallel MALA for Bayesian Logistic Regression
We next evaluated parallel MALA, focusing on (1) the convergence rate of Newton's 's method and the wall-clock time relative to sequential MALA, and (2) its efficiency in generating useful samples. For parallel MALA, used the stochastic quasi-DEER algorithm with I-sample estimate of the diagonal. We targeted the posterior of a Bayesian logistic regression (BLR) model of the German Credit Dataset with whitened covariates [60, 61] and (N I) prior. We compared 0.0015 (=80% acceptance rate) across batches of independent chains B 1, 2, 4, 8, 16, 32 and chain lengths L E 1K, 2K, 4K, 8K, 16K, 32K, 64K. Each configuration was repeated across 20 random seeds. The main text reports representative subsets of B and L highlighting key trends; full results and metric details are in Appendix C.I.
parallel and sequential MALA with step size
Wall-Clock Time and Convergence. From Figures 3A and 3C, we observe that for all tested sequence lengths L and batch sizes B except for B 32, parallel MALA can generate the same B batches of samples each much faster than sequential MALA, achieving speedups of up to 20 or 30 times for some smaller batch sizes. However, sequential MALA is slightly faster than parallel MALA for batch size B- = 32. This reflects a trade-off in the allocation of parallel resources, where for large batch sizes parallel MALA saturates GPU resources, degrading performance. Eventually, parallel MALA resulted in out-of-memory errors with B = 32 and L = 64K, and scaling to this size of samples or larger with parallel MALA would require the sliding window technique. Notably, parallel MALA typically converged in tens of Newton iterations (fig. 3B), and only rarely had chains not yet converged our prespecified
(red dashed line in Figure 3B and Appendix B).
On Sample Quality via MMD and Newly-Accessible Tradeofls. We next investigated the quality of samples generated via parallel MALA and whether parallel or sequential sampling was more efficient at generating high quality samples. To assess sample quality, we computed the Maximum Mean Discrepancy (MMD) [62] between sets of MALA samples and ground truth set of samples computed via NUTS (Hoffman et al. [19], see Appendix B). This allowed us to compare sample quality across batch sizes and chain length with lower MMDs value implying higher sample quality.
Parallel MALA produced samples of virtually identical quality to sequential sampling as measured by MMD (Figure 4), while achieving over an onder-of-magnituue speedup for smaller batch sizes.
7

---

For the largest batch size (B = 32), however, parallel MALA was slower than sequential MALA. To determine which combination of B, L. and sampling mode most efficiently generated useful samples. we computed interpolated timing estimates to identify the fastest configuration for simulating B chains of I. samples. The purple stars in Figure denote such interpolated settings: for instance, rather than one run with B = 32 and I. = 16K, four parallel MALA runs with B 8 and I. = 16K will likely achieve comparable sample quality in less time, while avoiding GPU oversaturation.
Across nearly all batch sizes and sequence lengths in Figure 4, sequentially-rum interpo- lated parallel MALA (purple stars) lies below and to the left of sequential MALA (red trian- gles), indicating faster wall-clock times with equal or better sample quality. This result highlights new tradeoff in GPU resource allocation of parallelizing across chain length versus batching independent chains. Our find- ings suggest that parallel MALA attains op- timal performance when more resources are
MMD vs. Wall-Clock Time Tradeoffs
88
par.
Seq.
B-32
Interp.
MMD (Lower is Better)
10-3
10-4
allocated for chain-length parallelization.
5.3 Parallel HMC
64K
We next evaluated the parallel HMC algo- rithms for simulating samples from the same BLR model of the German Credit dataset. First, we compared the iterations until conver- gence for parallelizing leapfrog integration using DEER, quasi-DEER, and block quasi- DEER with I, = 32 leapfrog steps across a range of step sizes and initial conditions (Fig- ure 5A). We found that DEER converged in the fewest iterations, although it incurs a very high memory cost for this problem. Quasi- DEER incurs a light memory cost, but in this example often required close to the maximum number of parallel iterations to converge (L iterations). Block quasi-DEER, which accounts for position and momentum interactions, notably reduced the number of iterations to convergence by Rs 2x compared to quasi-DEER across a range of relevant step sizes. In our timing experiments, we therefore solely considered block quasi-DEER for parallelizing the leapfrog integration.
10-2
10-1
10c
10a
Wall-Clock Time (sec)
Figure 4: MMD vs. wall-clock time performance for parallel and sequential MALA. Blue circles denote par- allel MALA, red triangles denote sequential MALA, and purple stars denote interpolated parallel MALA (e.g.. se- quentially running two independent instances instead of one B = 8, T. = 16K run). Fach line connects points equal batch size across varying sequence lengths L. Points closer the bottom-left indicate greater compute- time efficiency faster and higher quality).
= 4, T. = 16K
We compared the time to generate 1000 HMC samples from the BLR model across varying numbers of leapfrog steps and step sizes using sequential or parallel leapfrog integration (fig. 5B). We simulated a batch of four chains and estimated efficiency using effective sample size (ESS) per second using ArviZ [63]. For relatively larger numbers of leapfrog steps, parallel leapfrog integration achieved substantial efficiency gains over sequential leapfrog (fig. 5B), while for the largest step sizes and fewest leapfrog steps sequential HMC was often more efficient. We also compared peak ESS/s across hyperparameter settings for sequential HMC, HMC with parallel leapfrog, and HMC parallelized across the sequence length (each with four chains). In this setting where the optimal number of leapfrog steps is typically small, parallelizing HMC across the sequence achieved the highest FSS//.. While fewer leapfrog steps were optimal here, the efficiency gains from parallel leapfrog at larger leapfrog step counts may prove valuable in other problems. Appendix C further demonstrates parallel leapfrog integration for HMC on a -ileennniio item-response model [64].
5.4 Scaling Parallel MCMC Memory-Efficien quasi-DEER,
The stochastic quasi-DEase algorithm was critical for our MALA BLR experiments. For parallelizing MALA, Figure 6 compares its wall-clock performance against quasi-DEER implementation that uses automatic differentiation to compute the Jacobian diagonal, as the diagonal was not easily available in closed-form. The memory-efficient stochastic version achieved over 10x faster wall-clock times and remained in memory in multiple additional settings.
Early-Stopping. It is well-established that traditional MCMC algorithms inherently output im- perfect, approximate samples of the posterior. As such, reasonable question is whether we must
8

---

A
B
Convg. Iters., 32 Leapfrog Steps
Parallel/Seq. ESS/s
Peak ESS/s w/4 Chains
60
+00
0.005
4.00
250K
Iterations
0.01
2.00
#DEER
Speed
125K
20
0.05
0.50
0.1
0.25
0.00
0.02
0.04
0.06
16
24
40
96
Seq.
Stepsize
Leapfrog Steps
Leap
Seq.
Figure Parallel HMC sampling targeting BLR posterior. (A) Block quasi-DEER converges more efficiently than quasi-DEER for parallelizing leapfrog inlegration. (B) Relative ESS/s speedup of parallel vs.
sequential leapfrog. (C) Peak ESS/s across explored hyperparameter settings running 4 chains with sequential HMC, HMC with parallel leapfrog, parallel HMC across the sequence.
parallel MCMC until full convergence obtain useful samples for approximating posterior expectations, or we can get comparable quality samples in much faster time via carly-stopping at intermediate Newton iterations. To probe this question, we used parallel MALA as test case. Indeed, early-stopped parallel MALA can often yield samples with comparable MMD to running parallel MALA until convergence generating sequential samples (fig. 6B). For example, with B = 2 chains, running parallel MALA for only 8 iterations would have yielded samples with near comparable MMD quality to those at convergence, but also with nearly order of magnitude faster runtime. However, early-stopped parallel MALA performance varies across settings (Appendix C.I) and the optimal number of early-stopping iterations is likely problem-dependent. We leave more thorough investigation this to future work.
B
Wall-Clock Speedup of Mem-Efficient over AD Q-DEER
Early-Stop MMD vs. Wall-Clock
25
16K
64K
10-2
0.0 1-0-4-00
Par. (4)
32K
Par. (8)
Speedup
15
MMD
* Par. (All
10
10
10-1
B 1
8 4
B B 16
10-3
10-2
10-2
10
Batch Size (B)
Wall-Clock Time (Seconds)
Figure 6: Effectiveness of methods for scaling parallel MCMC: memory-efficiet quasi-DEER and early stopping. (A) Wall-clock speedup multiplicrs of memory-effiiient quasi-DEEs DEER over automatic diffcrentiation (AD) quasi-DEER (dotted line is equal runtime). Error bars denote 90% confidence intervals over 20 random seeds. Memory-efficient quaxi-DBi -EER consistently achieves over 10x 10x_faster faster wall-clock times. AD quasi-DEER out of memory for (B, L) [[4, 64K), (8, 32K), (16, 16K+)), so bars appear for these cases. (B) MMD-wall-Cee wall-clock tradeoffs for B Parallel MALA with/without early stopping. Points on the same line share the same Newton iteration count (4, 8, 12, full convergence); leftmost and rightmost points correspond to = 1K and L = 64K, respectively. Points closer to the bottom-left indicate greater compute-time efficiency. One can often achieve comparable MMD with far fewer Newton iterations than needed for full convergence.
5.5 Sentiment Classification from LLM Embeddings
We next scaled parallel MALA to a higher-dimensional sentiment classification task. We encoded reviews from the IMDB dataset [65] into 768-dimensional embeddings using 001 [66], partially inspired by Harrison et al. [67]. We targeted a BLR model with a ridge prior that predicted binary sentiment for 1024 randomly selected reviews given the embeddings. We simulated B chains with 4096 samples using sequential and parallel MALA with a step size of 0.015. To handle the larger dimensionality, we applied the sliding window method from Section 3.4 (considering window sizes of 128, 256, 512, and 1024) and used a precom- puted orthogonal basis transformation from the left singular vectors of the feature covariance matrix. Figure 7A depicts the sliding window approach and shows convergence of parallel samples to the sequential trace. We found that parallel MALA achieved the fastest runtime with a window size of 256 (Figure 7B) and achieved over 3x faster wall-clock time than sequential MALA (Figure 7B-D).

---

A
Parallel 100
Wall-Clock Times
C
Iters. Convg.
D
Wall-Clock Speedup
1st Coord. 1st Coord.
Seq.
Par.
0.4
Iterations (tz)
250
3
Time (sec)
Speedup
Parallel Iter 190
0.3
Par.
Seq.
200
2
2.5
150
0.2
0.0
100
1000
2000 3000
4000
27
2
2
2
2
2
2
2
27
2 29 2
Sample
Window Length
Window Length
Window Length
Figure 7: Sentiment classification from LLM embeddings. (A) Parallel and sequential samples of the first intermediate iteration (top) and convergence (below). The intermediate
coordinate of the posterior iteration shows the sliding window approach: the bluc-shaded area is the current window where the trace updated. (B) Wall-clock times for parallel and scquential MALA for sampling chains as a function window length. (C) Number of iterations to convergence for Parallel MALA as function of window length. (D)
Wall-clock speedup multiplier of parallel over sequential MALA for different window lengths.
6 Conclusion
In this paper, we introduced a novel framework for parallelizing MCMC across the sequence length using parallel Newton iterations. We demonstrated how several widely-used MCMC algorithms can be parallelized using this framework. In multiple experiments, showed this approach can generate many thousands of samples in tens of iterations, without sacrificing sample quality. In some cases, approach generated samples order of magnitude faster than traditional sequential methods. To obtain these results, we developed more efficient approaches for the parallelization of nonlinear recurrences, which is of independent interest. Additionally, this framework provides new opportunities and tradeoffs for the allocation of compute resources. First, one can now tradeoff parallelization of both the batch size and the sequence length. Second, for early-stopping MCMC we provide the ability to dynamically adapt the number of parallel iterations. Ultimately, this approach offers broadly applicable means of accelerating MCMC, a workhorse of modern statistics.
Limitations and Future Work. While our results demonstrate promising acceleration of MCMC via parallel Newton's 's method, several potential limitations remain. We found that convergence of Newton's 's method can vary with the sampler step size, target distribution geometry, and Newton- method hyperparameters. Our experiments also primarily targeted unimodal distributions; although we efficiently parallelized sampling of a mixture of Gaussians (Appendix B.5), highly-multimodee targets may have unstable local Jacobians, requiring damped updates alternative techniques [8]. Future work should establish a theory of when MCMC sampling is efficiently parallelizable and provide guiding principles for adaptive hyperparameter selection. clear opportunity is leverage results of contractivity in MCMC [68, 69) and to adapt concurrently-develppee theory on DEER convergence [70] to the MCMC setting.
Next, the presented algorithms trade increased memory and computational cost for reduced time complexity by evaluating many additional functions and gradients in parallel. For target distributions with high memory requirements, parallel MCMC may be less advantageous. Moreover, it remains unclear whether allocating parallel resources to more chains benefit: this is an open question warranting further theoretical and empirical study. In general, future work should scale these methods to nigher-dimensione and more challenging targets, and explore specialized hardware-aware acceleration strategies to further improve efficiency.
longer chains yields greater
Our HMC experiments used fixed step sizes and trajectory lengths, unlike adaptive HMC samplers such as NUTS [19]. Because NUTS introduces additional control flow to build sets candidate points, future work should examine its compatibility with the Jacobians required by DEER. However, simpler dynamic HMC variants that randomly jitter step sizes and leapfrog step counts should remain compatible with our methods. Beyond these HMC extensions, future directions also include applying this framework to stochastic gradient samplers [71-73] and recent MCMC algorithms [74, 75].
Finally, the proposed algorithms have several immediate applications. They can accelerate machine learning systems based on Langevin dynamics [76-78], enable faster approximate inference via early stopping, and improve training algorithms that rely on MCMC sampling, such as Monte Carlo EM and contrastive divergence [79]. In general, we envision this work unlocking MCMC sampling in domains where sequential sampling was previously too time-consuming.
10

---

Acknowledgments and Disclosure of Funding
We thank Art Owen, Kelly Buchanan, and members of the Linderman Lab for helpful feedback. This work was supported by grants from the NIH BRAIN Initiative (U19NS1 1 13201, ROINS1 & RFIMH133778) and the NSFNNH CRCNS Program (ROINSI30)) 137789). W. would like to acknowl- edge support from an NSF Graduate Research Fellowship. X.G. would also like to acknowledge support from the Walter Byers Graduate Scholarship from the NCAA. L.K. was a Goldstine Fellow at IBM Research while working this paper. S. S.W.L. is supported by fellowships from the Simons Collaboration on the Global Brain, the Alfred P. Sloan Foundation, and the McKnight Foundation. The authors have competing interests to declare.
Some of the experiments were performed on the Stanford Sherlock computing cluster. We thank Stanford University and the Stanford Research Computing Center for providing computational resources and support that contributed these research results.
References
[1] Persi Diaconis. The Markov chain Monte Carlo revolution. Bulletin of the American Mathemat- ical Society, 46(2): 179-20,, April 2009,
[2] Charles Geyer. Introduction to Markov chain Monte Carlo. Handbook of Markov chain Monte Carlo, 201 16022(45):22, 2011.
[3] Andrew Gelman, John B. Carlin, Hal S. Stern, David B. Dunson, Aki Vehtari, and Donald B. Rubin. Bayesian Data Analysis. CreateSpace, United States, 3rd ed edition, 2013. ISBN 978 178143984
[4] Pavel Sountsov, Colin Carroll, and Matthew D Hoffman. Running Markov Chain Monte Carlo on Modern Hardware and Software. arXiv preprint arXiv:2411.04200 2024.
[5] Charles C Margossian, Matthew Hoffman, Pavel Sountsov, Lionel Riou-Durand, Aki Vehtari, and Andrew Gelman. Nested Rhat: Assessing the convergence of Markov chain Monte Carlo when running many short chains. Bayesian Analysis, 1(1):1-28, 2024.
[6] Federico Danieli, Miguel Sarabia, Xavier Suau Cuadros, Pau Rodriguez. and Luca Zappella. Deepper: Parallelizing sequential operations in neural networks. Advances in Neural Information Processing Systems,
2023.
[7] Yi Heng Lim, Qi Zhu, Joshua Selfridge, and Muhammad Firmansyah Kasim. Parallelizing non-linear sequential models over the sequence length. In International Conference on Learning Representations, 2024.
[8] Xavier Gonzalez, Andrew Warrington, Jimmy T. H. Smith, and Scott W. Linderman. Towards scalable and stable parallelization of nonlinear RNNs. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
[9] Stuart Geman and Donald Geman. Stochastic relaxation, Gibbs distributions, and the Bayesian restoration of images. IEEE Transactions on Pattern Analysis and Machine Intelligence, PAMI-6 (6):721-741, 1984.
[10] Julian Besag. Comment on "Representations of knowledge in complex systems" by Grenander and Miller. Journal of the Royal Statistical Society: Series B (Methodological), 56(4):549-58) 1994.
[11] Simon Duane, Anthony D Kennedy, Brian Pendleton, and Duncan Roweth. Hybrid Monte Carlo. Physics Letters B, 195(2):216-222, 1987.
[12] Radford M Neal et al. MCMC using Hamiltonian dynamics. Handbook of Markov Chain Monte Carlo, 2(11):2, 2011.
[13] Michael Betancourt, conceptual introduction to Hamiltonian Monte Carlo. arXiv preprint 1xX1v1100.00434,200 2017. 2017.

---

[14] Guy E. Blelloch. Prefix sums and their applications. Technical Report CMU-CS-90-190, School of Computer Science, Carnegie Mellon University, November 1990.
[15] Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. arXiv preprini arXiv: arXiv: 04057, 2017.
[16] Jimmy T.H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling. In International Conference Learning Representations (ICLR), 2023.
[17] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint urXiv:2312.00752, 2023.
[18] The JAX Authors. Advanced Automatic Differentiaiion, 2024.
[19] Matthew D Hoffiman, Andrew Gelman, et al. The No-U-Turn sampler: Adaptively setting path lengths Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(1):1993-1623, 2014.
[20] Costas Bekas, Effrosyni Kokiopoulou, and Yousef Saad. An estimator for the diagonal of matrix. Applied Numerical Mathematics, 57(11-12:1214-1229, 2007.
[21] Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney. Adahessian: An adaptive second order optimizer for machine learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 10665-10674 2021.
[22] Michael Hutchinson, A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications in Paatitics-Simullttie and Computation, 18(3): 1059 1076, 1989.
[23] Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sampling of diffusion models. Advances in Neural Information Processing Systems, 36:4263 4276, 2023.
[24] Mónika Farsang, Ramin Hasani, Daniela Rus, and Radu Grosu. Scaling Up Siquid-Resistann Sqquid-Caaccitance Networks for Efficient Sequence Modeling. In Advances Neural Infor- mation Processing Systems (NeurIPS), 2025.
[25] Yvon Maday and Gabriel Turinici. A parareal in time procedure for the control of partial differential equations. Comptes Rendus Mathematique, 335(4).387-32, 2002.
[26] Nikil Selvam, Amil Merchant, and Stefano Ermon. Self-refining diffusion samplers: Enabling parallclization via parareal iterations. Advances Neural Information Processing Systems, 37: 5429-5453, 2024.
[27] Zhiwei Tang, Jiasheng Tang, Hao Luo, Fan Wang, and Tsung-Hui Chang. Accelerating parallel sampling of diffusion models. In Forty-first International Conference Machine Learning, 2024.
[28) Xavier Gonzalez, E. Kelly Buchanan, Hyun Dong Lee, Jerry Weihong Liu, Ke Alexander Wang. David M. Zoltowski, Christopher Ré, and Scott W. Linderman. Unifying Framework for Parallelizing Sequential Models with Linear Dynamical Systems. arxiv/2509.21716, 2025.
[29] Nima Anari, Sinho Chewi, and Thuy-Duong Vuong. Fast parallel sampling under isoperimetry. In The Thirty Seventh Annual Conference on Learning Theory, pages 161-185. PMLR, 2024.
[30] Elaine Angelino, Matthew James Johnson, Ryan P Adams, et al. Pattems of scalable Bayesian inference. Foundations and Trends@ in Machine Learning, 9(2-)): 19921999277222 2016.
[31] Sebastiano Grazzi and Giacomo Zanella. Parallel computations for metropolis markov chains with picard maps. arXiv preprint arXiv:2506.09762, 2025.
[32] Bob Carpenter, Andrew Gelman, Matthew D Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A probabilistic programming language. Journal of Statistical Software, 76:1-32, 2017.
12

---

[33] Du Phan, Neeraj Pradhan, and Martin Jankowiak. Composable effects for flexible and acceler- ated probabilistic programming in NumPyro. arXiv preprint arXiv:1912.1155,, 2019.
[34] Junpeng Lao, Christopher Suter, Ian Langmore, Cyril Chimisov, Ashish Saxena, Pavel Sountsov, Dave Moore, Rif A Saurous, Matthew D Hoffman, and Joshua V Dillon. tfp.meme: Modern Markov chain Monte Carlo tools built for modern hardware. arXiv preprint arXiv:2002.01184 2020.
[35] Alberto Cabezas, Adrien Corenñlos, Junpeng Lao, Rémi Louf, Antoine Carnec, Kaustubh Chaudhari, Reuben Cohn-Gordon, Jeremie Coullon, Wei Deng. Sam Duffield, et al. Blackjax: Composable Bayesian inference in JAX. arXiv preprint wwXiv2240..10790 2024.
[36] Andrew Gelman and Donald B Rubin. Inference from iterative simulation using multiple sequences. Statistical Science, 7(4):457-472, 1992.
[37] Matthew Hoffman, Alexey Radul, and Pavel Sountsov. An adaptiveMMCMC scheme for setting trajectory lengths in Hamiltonian Monte Carlo. In International Conference on Artificial Intelligence and Statistics, pages 3907-3915. PMIRR 2021.
[38] Hugh Dance, Pierre Glaser, Peter Orbanz, and Ryan Adams. Efficiently vectorized MCMC on modern accelerators. arXiv preprini arXiv:2503. / /7405, 2025.
[39] Robert Nishihara, lain Murray, and Ryan P Adams. Parallel MCMC with generalized elliptical slice sampling. The Journal of Machine Learning Research, 15(1):0887-2122, 2014.
[40] Matthew D Hoffman and Pavel Sountsov. Tuning-free generalized Hamiltonian Monte Carlo. In International conference on Artificial Intelligence and Statistics, pages 7799-7813. PMLR,
2022.
[41] David Newman, Padhraic Smyth, Max Welling, and Arthur Asuncion. Distributed inference for Datent Dirichlet Allocation. Advances in Neural Information Processing Systems, 20, 2007.
[42] Joseph Gonzalez, Yucheng Low, Arthur Gretton, and Carlos Guestrin. Parallel Gibbs sampling: From colored fields to thin junction trees. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pages 324-332. JMLR Workshop and Conference Proceedings, 2011.
[43] Matthew J Johnson, James Saunderson, and Alan Willsky. Analyzing hogwild parallel Gaussian Gibbs sampling. Advances in Neural Information Processing Systems, 26, 2013.
[44] Jun S Liu, Faming Liang, and Wing Hung Wong. The multiple-try method and local optimization in Metropolis sampling. Journal of the American Statistical Association, 95(449):121-13 2000.
[45] Radford M Neal. Markov chain sampling for non-linear state space models using embedded hidden Markov models. arXiv preprint math/0305039, 2003.
[46] Hakon Tjelmeland. Using all Metropolis-Hastings proposals to estimate mean values. Technical report, 2004.
[47] Daan Frenkel. Speed-up of Monte Carlo simulations by sampling of rejected states. Proceedings of the National Academy of Sciences, 101(51): 2004.
[48] Ben Calderhead. A general construction for parallelizing Hatropolis-Hasttnng algorithms. Proceedings of the National Acudemy of Sciences, 1 1408811348 2014.
[49] Nathan E Glatt-Holtz, Andrew J Holbrook, Justin A Krometis, and Cecilia F Mondaini. Parallel MCMC algorithms: theoretical foundations, algorithm design, case studies. Transactions of Mathematics and Its Applications, 8(2)xtaae004, 2024.
[50] Anthony E Brockwell. Parallel Markov chain Monte Carlo simulation by pre-fetching. Journal of Computational and Graphical Statistics, 15(1):246-261, 2006.
[51] Flaine Angelino, Eddie Kohler, Amos Waterland, Margo Seltzer, and Ryan P Adams. Acceler- ating MCMC via parallel predictive prefetching. arXiv preprint arXiv:1403.7263 2014.
13

---

[52] Simo Särkkä and Ángel F
Temporal parallelization of bayesian smoothers. IEEE Transactions on Automatic Control, 66(1):299-336, 2020.
[53] Adrien Corenflos, Nicolas Chopin, and Simo Särkkä. De-sequentialized monte carlo: a parallel- in-time particle smoother. Journal of Machine Learning Research, 23(283):1-39, 2022.
[54] Fatemeh Yaghoobi, Adrien Corenflos, Sakira Hassan, and Simo Särkkä. Parallel square-root statistical linear regression for inference in nonlinear state space models. SIAM Journal on Scientific Computing, 47(2):4454 (4102))4554-47,,,, 2025.
[55] Zaijing Huang and Andrew Gelman. Sampling for Bayesian computation with large datasets. Available SSRN 1010107, 2005.
[56] Willie Neiswanger, Chong Wang, and Eric Xing. Asymptotically exact, embarrassingly parallel MCMC. arXiv preprint arXiv: 1311,4780. 2013.
[57] Steven L. Scott, Alexander W. Blocker, Fernando V. Bonassi, Hugh A. Chipman, Edward I. George, and Robert E. McCulloch. Bayes and big data: the consensus Monte Carlo algorithm. International Journal of Management Science and Engineering Management, 1 11())78888, 2016.
[58] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Bonderrman-Minne and Qiao Zhang. JAX: composable transformatoons of Python+NumPy programs, 2018. URL http:// bttt://ggthule wttp:/ccoo/////e/aaaaaaa
[59] Donald B Rubin. Estimation in parallel randomized experiments. Journal of Educational Statistics, 6(4):377-401, 1981.
[60] Dheeru Dua, Casey GrafT, et al. UCI machine learning repository. 7(1):62, 2017. URL http://archive. ics. wuw.uuuiuuu..u edu/ml.
[61] TFDS. TensorFlow Datasets, a collection of ready-to-use datasets. org/datasets, 2024.
[62] Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(25):723-773, 2012.
[63] Ravin Kumar, Colin Carroll, Ari Hartikainen, and Osvaldo Martin. ArviZ a unified library for exploratory analysis of bayesian models in python. Journal of Open Source Software, 4(33): 1 1143, 2019.
[64] Andrew Gelman and Jennifer Hill. Data analysis using regression and multileve/hhierarcicaal models. Cambridge University Press, 2007.
[65] Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies. pages 142-150, 2011.
[66] Jinhyuk Lee, Feiyang Chen, Sahil Dua, Daniel Cer, Madhuri Shanbhogue, Iftekhar Naim, Gustavo Hernández Ábrego. Zhe Li, Kaifeng Chen, Henrique Schechter Vera, et al. Gemini embedding: Generalizable embeddings from Gemini. arXiv preprint arXiv:2503.0789 2025.
[67] James Harrison, John Willes, and Jasper Snock. Variational Bayesian last layers. arXiv preprint arXiv:2404.1555 15599, 2024.
[68] Persi Diaconis and David Freedman. Iterated random functions. SIAM review, 41(1):45-76. 1999.
[69] Nawaf RabeR, Andreas Eberle, and Raphael Zimmer. Coupling and convergence for hamiltonian monte carlo. The Annals of applied probability, 30(3): 2020.
14

---

[70] Xavier Gonzalez, Leo Kozachkov, David M. Zoltowski, Kenneth L. Clarkson, and Scott W. Linderman. Predictability enables parallelization of nonlinear state space models. In Neural Information Processing Systems (NeurIPS), 2025.
[71] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on Machine Learning (ICML-I1), pages 681-688, Citeseer, 2011.
[72] Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient Hamiltonian Monte Carlo. In International Conference on Machine Learning, pages 1683-1691, PMLR, 2014.
[73] Ruqi Zhang, A Feder Cooper, and Christopher De Sa. AMAGOLD: Amortized Metropolis adjustment for efficient stochastic gradient MCMC. In International Conference on Artificial Intelligence and Statistics, pages 2142 2152. PMLR, 2020.
[74] Lionel Riou-Durand and Jure Vogrinc. Metropolis adjusted Langevin trajectories: A robust alternative to Hamiltonian Monte Carlo. arXiv preprint arXiv:2202. 13230, 2022.
[75] Lionel Riou-Durand, Pavel Sountsov, Jure Vogrine, Charles Margossian, and Sam Power. Adaptive tuning for Metropolis adjusted Langevin trajectories. In International Conference Artificial Intelligence and Statistics, pages $102-8116. PMLR, 2023.
[76] Will Grathwohl, Kuan-Chieh Wang, Jörn-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly energy based model and you should treat it like one. arXiv preprint (9/2vv29/22032000 2019.
[77] Erik Nijkamp, Mitch Hill, Song- Chun Zhu, and Ying Nian Wu. Learning non-convergent non-persistent short-run MCMC toward energy- based model. Advances in Neural Information Processing Systems, 32, 2019.
[78] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[79] Geoffrey Hinton. Training products of experts by minimizing contrastive divergence. Neural Computation, 14(8):1771-1800, 2002.
[80] Ernst Hairer, Gerhard Wanner, and Syvert P Nørsett. Solving ordinary differential equations I: Nonstiff problems. ponslifprroblems. Springer, 1993.
[81] Joshua V Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A Saurous. Tensorflow distributions, arXiv preprint arXiv: /71 I. 10604, 2017.
[82] Alexandros Beskos, Natesh Pillai, Gareth Roberts, Jesus-Maria Sanz-Serna, and Andrew Stuart. Optimal tuning of the hybrid Monte Carlo algorithm. Bernoulli, 19(5A):1511-154, 2013.
ISSN 13507265.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Please see our abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Please see our "Results" and "Conclusion" sections.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
16

---

Justification: Our paper is mainly computational and methodological, with less of a focus on theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Please see our "Parallelizing MCMC" and "Results" sections, as well as Appendices A and B. We also include full source code to reproduce all main experimental results.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
17

---

reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: please see our full source code in the supplementary materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /nips. wccppbllcc/guuueeeo
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Please see our "Parallelizing MCMC" and "Results" sections, as well as Appendices A and B. also include full source code with ample comments reproduce all main experimental results in paper in our supplementary materials.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All figures, whenever feasible, contain either error bars or confidence bands indicate statistical significance. We also explain in detail how these bars and bands computed.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
18

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: Please see our "Results" section, Appendices A and B, and our supplementary code.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https : http///eerrrp.eeeee wwweerris..cpppbllllee cc/ppbbicc/tticccuuuueeeeacceee
Answer: [Yes]
Justification: We have carefully read the NeurlPS Code of Ethics and have made sure that our research conforms to it.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
Justification: This paper introduces a framework of foundational research methods that does not have any direct paths to negative societal impacts, aside from the general interactions of machine learning and statistical methods with society.
19

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper does not involve any data or models that have high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the main packages (e.g.. JAX and TensorPlow Probability) and acknowledge the original creators of the DEER quasi-DEEE algorithms.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
20

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswithcom thcode. com/coo/////ooooo comddataset has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Please see the README in our supplemental code. Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooel as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The findings in this paper do not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The findings in this paper do not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
21

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core methods introduced this paper do not involve LLMS as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https: / httplpp.ccc/oooffeeeecccmc cc/Coonferencee/222//L... for what should or should not be described.
22

---

A Additional Algorithm Details
A.1 DEER Algorithm
The high-level steps the DEER and quasi-DEEs algorithms are shown in Algorithm 1, where *DEER is either the DEER quasi-DEEL Newton update.
Algorithm 1 Parallel evaluation of nonlinear sequence models using variants of DEER
Require: initial state Sp, function sequence f1.7, initial guess 11. tolerance 6, update +DEER
Amax f x
i 0
while Amax do
(f+1)
+DEER 1.7" f1.1, Sp) 5 (6)
> global update with Olog T) parallel time complexity
1:T
Aaax max; | (1+1)
8 + i+1
end while return 11.
A.2 HMC algorithm with parallel leapfrog
For clarity, the sequential HMC algorithm with parallel leapfrog evaluation is written in full in Algorithm 2. The inner loop of steps of leapfrog integration is evaluated using parallel Newton's method. It is also straightforward add diagonal mass matrix.
Algorithm 2 HMC Step with Parallel Leapfrog function HMC(p(x), (9) L)
Xo f- (() Vo + NNO,
sample momentum compute energy half-step momentum & evaluate in parallel
Hg + 2 Vo (xg)(xo)
Vo + Vo á Vx log p(xo)
for t + ++++a do
X; f- x;-1 + -
Vi Vt-1 xx log (xx))
end for
VL VL 2 Vx log p(x)
reverse half-step momentum compute energy
HL 2 VL V1 log p(XL)
U ~ Unif(0.1
if < min(1. (,, H Ho) then
Accept-Reject Step
return (1+1)
else
return x(0
end if
end function
A.3 Block quasi-DEER for HMC with Parallel Leapfrog
Here we discuss parallelizing leapfrog with block quasi-DEER in additional detail. Recall that for HMC with parallel leapfrog, the state the function being parallelized the concatenation of position and momentum states s, [x, v.] with E R²d The Jacobian of the leapfrog step used in HMC with parallel leapfrog integration has a block structure:
0f Øs
J(s
(s,-1)=
8x--1 9v.
(12)
dv.
Fx.
1
ID 22 log 1og
ED Id + Vo vvv-1)
(13)
23

---

where Dp is the identity matrix and V 1ogp(xi-1-ev--- is the Hessian of log p(x) evaluated at the new position state x; x + -. We note that here we have derived the Jacobian in the absence of a mass matrix, but it is straightforward to add such matrix and we use diagonal mass matrix in the stem-response theory experiment. The top two blocks of this Jacobian are diagonal. While the lower blocks depend on the target distribution, the lower-right block is also equal to diagonal matrix plus the Hessian term scaled by e² such that it will be close to diagonal for small step sizes. The structure of this Jacobian motivated the block quasi-DEER approximate Jacobian, which retains only the diagonal of each block
diag( 8xx Dv:
8x. diag( itv: diag(
J BQQ
(14)
diag( ox.-1
div Dva-1
For parallelizing leapfrog integration in HMC, this approximate Jacobian is
ID J JBQ(SL-1)= = ediag v p(0 evl-1)) Id + e² diag log p(x + ev--)))
dD
(15)
Importantly, we are efficiently able to construct single-sample stochastic estimate of the block quasi-DEER Jacobian using only one Hessian-vector product since
diag( V pog ev.-1)) [ad
(V logp(x;,1 + +v=-1)
(16)
The Hessian- vector product can be efficiently computed using a combination of reverse-mode mode and forward-mode autodifferentiation.
An important feature of quasi-DEER, which retains only the diagonal of the Jacobian, is it can be implemented using a parallel linear recursion with O(TD) memory and O(TD) work. Here we show how block-quasi DEER also admits such recursion with constant additional factor of memory and work. Consider two matrices A and B that are each two-by-two block matrices and each block is diagonal with
a1
b
bp
ddiag(a) diag(c)
diag (b)
A =
C1
di
=
diag(d)
(17)
CD
dp
and
diag(e) e) diag( g
diag(f)
B
diag(h).
(18)
The matrix product of A and B is
diag a O e b
diag(a 3 f 6 diag C C do O diagic of d
AB
(19)
Notably, the product of A and B is also a block matrix composed of diagonal blocks and this product be computed with element-wise products of size D. When the matrix A is of size 2D x 2D, the linear recursion requires OT44D) memory to store the diagonal and two off-diagonals for each time point and OTTD) work compared O(T2D) memory and O(T2D) work for linear recursion of purely diagonal matrices. The resulting linear complexity in D dramatically improves over the cubic complexity in D for DEER when the dimension of the system is large. We hypothesize that this parallel linear recursion used for block DEER could be useful in other applications of parallelizing nonlinear dynamics with interacting variables and in deep state space models. Additionally, while we constructed block quasi-DEER with 2 x 2 block matrix for this problem, we also note that a similar scheme possible for block matrices that consist of x blocks with each block being diagonal. We note that this approach has OTT: memory and O D) work, and so provides natural interpolation between standard quasi-DEER, with (T D) memory and OTDD) work. and standard DEER, with OTTD²) memory and OTT)3) work.
24

---

B Additional Experiment Details
In our applications we use combination of absolute and relative tolerance [80] the stopping criterion for parallel Newton's 's method. In all experiments, we set the absolute tolerance to 10 and relative tolerance to 10-3 unless otherwise noted.
B.1 Parallel MALA
We use Bayesian logistic regression (BLR) on the German Credit Dataset (with whitened covari- ates) [60, 61] our testhed for this of experiments. Specifically, will use guasi-DEER-basse
Parallel MALA to generate B independent chains of I. samples each from the posterior distribution of the BLR model with a (0, I) prior the covariates. We vary the batch size B E (1, 2, 4, 8, 16, 32) and the sequence length L E {1000, 2000, 4000, 8000, 16000, 32000, 640000. To solidify confidence in our results, we will repeat all settings across 20 random seeds. We repeat the corresponding set- tings for the baseline sequential MALA algorithm. To ensure fairness, we implement the sequential procedure using the JIT-compilable jax. lax. scan module.
For both parallel and sequential MALA, we use a learning rate of e = 0.0015, selected to maintain roughly x 80% acceptance ratio. For both algorithms, we sample independent initial states from their (( I) priors for each of our B chains. Then, we run w = 3 burn-in steps to orient these initial states into more cetisticalll-llussibbe regions of the parameter space. We then set our initial guess of states for all timepoints at Newton iteration 0 to the initial state such that $1.7 = Sp.
For parallel MALA, we set tter of 50 + 5L x 10-1, chosen heuristically because expect longer sequence lengths L to require more iterations until all sequence elements numerically converge. We set the absolute tolerance to 10-4, 4 From Appendix Figures 11 and 15, we see that the vast majority of our trials across settings reach numerical convergence before their corresponding max_iter iter values. This suggests that our max mar_iter cutoffs were appropriate.
To compute wall-clock times, for each random seed, we first perform 3 warm-up runs of each evaluated algorithm. Then, we perform 5 timing runs of each evaluated algorithm, logging the total wall-clock time across the 5 timing runs. Finally, we divide this total time by 5 to return our wall-clock time for this seed. The warm-up runs serve account for potential slowdowns (of both the sequential and parallel MALA) from JIT-compilation overhead on initial executions. We average over 5 timing runs to reduce variance from uncontrollable system processes. We do note, qualitatively, that the one-time cost of JIT-compilation did seem to be noticeably higher for our parallel algorithms compared the sequential baselines. However, in typical use cases involving repeated sampling (i.e. repeated sampling calls), runtime over many calls will almost certainly dominate compile time. As such, we believe that our timing methodology well-justified. Nonetheless, reducing this initial overhead remains a useful direction of future work.
For all experiments measuring wall-clock time, we set full_fuuaceeeleeeel trace=False for quasi-DEER driver so that we only store the output after the last Gauss-Newton iteration. Each run of quasi-DEER under this setup also outputs the numbers of iterations required for convergence of each of our B chains, allowing us to generate Appendix Figure 11. To generate Figures 14 and 15, we require access to the quasi-DEER samples after each Newton iteration, not just after convergence or hitting the maximum iteration count. As such, for these experiments, we set full_ trace=True.
All experiments were run NVIDIA H100 GPU with 80 GB of RAM.
B.1.1 Measuring Sample Quality via Maximum Mean Discrepancy (MMD)
To assess sample quality, we use Maximum Mean Discrepancy [62]. We do not use traditional MCMC metrics like effective sample size (ESS) or R because such metries mainly target the autocorrelatednes of traditional sequential MCMC algorithms. In contrast, our parallel MALA algorithm by nature does not proceed sequentially sample-by-sampe and the autocorrelation of its generated samples are inherently limited by the autocorrelation of the base sequential MALA algorithm.
Given that the dominant use case for generating MCMC samples is to use said samples towards approximating the posterior expectation of some target function of interest, we can quantify sample quality by how closely our sample- based estimates can match the true target posterior expectations.
25

---

Formally, for probability distributions P2 and function class F, the (theoretical) Maximum Mean Discrepancy (MMD) is defined by Gretton al. [62] as MMD(F, Pr. P2) = sup -
FEF
for X ~ P1 and Y ~ P2- we set Pa be the true target posterior distribution and let p1 represent the distribution represented by our generated samples, then MMD is a natural measure of sample quality.
However, we do not have access to the true target posterior distribution p2- Thus, we approximate the "ground truth" P2 via large set of gold standard MCMC samples generated using the state-of-the-art No U-Turn Sampler with Dual-Averaging Stepsize Adaptation introduced by [19] and implemented in TensorFlow Probability [81]. Specifically, following standard best practices, we aim for target acceptance probability of 75% with initial states being drawn from the A I) prior. To
ensure high quality, we perform 50, 000 burn-in steps before collecting 100, 000 samples our gold standard approximation of the true target posterior.
Yet, taking the supremum over all in intractable. As such, following Gretton et al. [62], we pick a function class F corresponding to a reproducing kernel in our case, we will use Gaussian RBF kerne::
connntially-uccuuntabbbl function class F is analytically
|| - x21
k( x22 == exp
20²
with hyperparameter a to be discussed in the next paragraph. Let {a: refer to the samples outputted by sequential or parallel MALA, with N B and let ر represent the gold standard samples. In such setting, we can empirically and unbiasedly estimate the MMD of our samples with respect to the true target posterior as follows (see [62]):
N
MMD(FR = N(N ==1
M
M
M N
2
kV,,
k( v,)-
M(M
MN
Because we are working with high-dimensional settings with large batch sizes B and sequence length L, we cannot store all of our samples in GPU memory once. As such, we approximate each term above using at most Mmax = 50, 000 randomly-selected samples from our MALA outputs and take the average of our MMD estimates over 10 replications.
Finally, following Gretton et al. [62], we select a to be the median pairwise distance between samples in our gold standard set. Because storing the entire pairwise distance matrix would overwhelm our GPU memory, we randomly select 50, 000 samples from this set, compute the median pairwise distance of this subset, and then take the mean over 10, 000 repetitions of this procedure.
B.2 Gibbs sampling experiment
The generative model for the Gibbs sampling experiment is
7² x -20, 7
# NN 7² /²//Ko)
(20)
0. ~ NN,N
03 ~ x (00,03)
Fan. ~ NO,, 03)
(21)
where V/ = 0.1, TO = 100.0, Po 0.0, No = 0.1, on = 0.1, and of 10.0. We paral- lelized Gibbs sampling coordinate updates that simulate samples from the 18-dimensional posterior p(72, ur (0., - ==1 X. Φ) where X are the set of observations across the S = 8 schools each with N 20 students and the hyperparameters are Φ = {Mg, To, /g: Ko; Ogg, ) We simulated data such
that the means and standard errors corresponded to those in Table 5.2 of Gelman and Hill [64]. The
dynamics function at time-step proceeded via the following coordinate updates
-² p(7? 14-1, 10
(22)
#t p(u| (0,,1-1) -2
(23)
0,t p(a,,t
Meg -2 02 X,
(24)
~ ~,,~plo' 00,,, Φ)
(25)
26

---

where 8,,, denotes the sampled value for 6, at time point t. We reparameterized this update using standard reparameterization for Gaussian random variables and the fact that if 2 x ( 7²) then kz x 2 [v, R)²) for the 7² and o² given that the location parameter of the x x- distributions were fixed and known for each of those variables. In our implementation, we passed random number key as input each step of the function and sampled random variates as part of ft.
For each chain, we first simulated a draw from the prior distribution and then ran one sequential sampling step targeting the posterior. The result of this sequential sampling step was to be the initial state Sn for each chain. For Newton's 's method, the guess of the state sequence at iteration 0 was set to the initial state such that s So for all & .,.... T. using quasi-DEER for this problem, we used 3-sample stochastic estimate of the diagonal. We additionally applied tuned diagonal preconditioner to the stochastic quasi DEER Jacobian estimates.
as in the MALA experiment. When
To compare the timing of sequential and parallel sampling, we ran chains of batch size B {1, 4,8 8, 16, 32, 64) and length T. E (8000, 16000, 32000, 64000, 128000). For all batch sizes 64, we additionally simulated chains of length E (256000, 512000, 1024000). On CPU, we ran sequential sampling for all batch sizes and for lengths up to T. 32000; for values of > 32000 we extrapolated the time for sequential sampling (i.e. for = 64000 we doubled the amount of time it took to simulate L 32000 samples). We repeated each run 5 times and computed the timing a similar manner as our MALA experiments. That is, we computed the average time to run the JTT-compiled code across 3 timing runs after performing 2 warm-up runs of each algorithm.
except for
35
0.20
sasmpless
29
0.15
11X,0
00
0. 10
0.05
2000
6006
8000
1003
10000
0.00
10
15
20
25
30
35
Iteration
Parale Olbbs
Sequental Gbos
Figure 8: Example samples from sequential and parallel Gibbs sampling. Left: Samples across iterations for sequential and parallel Gibbs sampling. The two sample traces are overlapping. We note that we transformed 7² into т for visualization purposes. Right: The corresponding histograms of samples of p(r X, Φ) are identical.
B.3 Sequential HMC with parallel leapfrog experimental details
We also applied sequential HMC with parallel leapfrog to the Bayesian logistic regression model of the German -Ceddit dataset [60, 61] with whitened covariates. We focused on sampling 4 batched chains and swept across the number of leapfrog steps
{4, 8, 12, 16, 20, 24, 32, 40, 64, 96) and step sizes {0.005, 0.0075, 0.01, 0.02, 0.03, 0.04, 0.05, 0.075, 0.1). Step sizes at 0.1 and above and e, we
resulted unstable integration and acceptance rates near zero. For each combination estimated the time to simulate 2000 samples from HMC using either sequential or parallel leapfrog integration using a JIT-compiled function. We repeated each step size and number of leapfrog steps across 10 random seeds. In all cases, only used the absolute tolerance. For each run, computed the average effective sample size across parameters. We then reported the average effective sample size per second. A relative comparison of the average ESS/s across runs between HMC with parallel or sequential leapfrog is reported in Figure 5B. Additionally, the maximum ESS/s across settings of L and for HMC with parallel or sequential leapfrog is shown in Figure 5C.
When parallelizing the leapfrog algorithm, the initial state the concatenation of the initial position and the momentum after a half-step So [xo, vo). We set the guess of the sequence of states to be equal to this initial state, as in the Gibbs and MALA setups. Across different initial states, compared the number of iterations until convergence for DEER, block quasi-DEER and quasi-DEER
27

---

across the same set of step sizes and with L = 32 steps for BLR. These results are reported in Figure 5A.
We next explored using sequential HMC with parallel leapfrog for an item-repsonse theory model [64, 61]. The target distribution in this case has 1> 501 dimensions and the dataset has approximately 30K datapoints, making this substantially larger problem. We compared the time to run HMC with either sequential or parallel leapfrog across 1 or 2 chains and varying step sizes and number of leapfrog steps. For each setting we ran 10 repeats across different random number seeds. We estimated the time to simulate 1000 samples from the posterior distribution (Figure 19). Here, we used pre-tuned diagonal mass matrix, which empirically we found helped improve the convergence of parallel leapfrog integration. We found that when simulating 1 chain the HMC with parallel leapfrog can achieve speed-ups over HMC with sequential leapfrog. However, at two chains the sequential leapfrog was faster.
B.3.1 Step-size for HMC with parallel leapfrog
Tuning the step size and number of leapfrog in HMC is important for optimal exploration of the target distribution. The standard analyses for determining the optimal step size and acceptance rate typically assume a linear time and computational cost in the number of leapfrog steps [12, 82]. However, parallelizing the leapfrog integration enable sublinear cost in the number of steps, violating this assumption. A second compounding factor is that the number of parallel Newton iterations until convergence of the parallel integration can vary function of the step size, as saw in results. New analysis investigating the appropriate optimal step sizes and acceptance rates need to be devised under these conditions. We hypothesize that in some cases, these conditions will favor the use of smaller step sizes and more leapfrog steps, thereby increasing the optimal acceptance rate. We currently leave a formal analysis of this future work.
B.4 Parallel HMC across the sequence
We parallelized HMC across the sequence for both the Bayesian logistic regression example and for the Rosenbrock distribution. For parallel HMC on Bayesian logistic regression we used memory- efficient quasi-DEER to parallelize the algorithm. We tuned the step size, number of leapfrog steps, and hyperparameters of the algorithm to consistently achieve fast convergence with high numbers of samples. The resulting parameters were L = 4 leapfrog steps and step size 0.04. We simulated five sequential samples before setting the initial state for parallel HMC. In this case, we simulated batches of chains of length 1000. The initial guess of the states for parallel Newton's method was set to the initial state for all time points. When using memory-efficient quasi-DEER, we clipped the magnitude of the diagonal Jacobian be less than 1.0 and we used a damping factor of 0.4 [8]. We repeated the sampler across 10 different seeds. The sampler typically converged in around 20 iterations and the acceptance ratio was around 0.89%. We computed the ESS/s in the same manner as for sequential HMC with sequential or parallel leapfrog in the previous subsection.
In Figure 1, we parallelized HMC sampling of the Rosenbrock "banana" distribution implemented in TFDS [61]. We parallelized a chain of 100K samples where each HMC consisted of L = 8 leapfrog steps with step size = 0.3. This resulted in an acceptance rate of approximately 97 97.6%. For this 2D distribution, we used the full DEER Jacobian scaled by damping factor of 0.5 [8] and clipped to be within [-1, 1]. For this motivating example, we ran this with a single random seed and did not examine timing or convergence profiles across runs.
B.5 Multimodal mixture of Gaussians example
Here we demonstrate that the approach can efficiently parallelize sampling from a multimodal distribution. We constructed mixture-of-Gaussian distribution in two dimensions with four components (Figure 9). We used MALA with step size of c 0.1 to simulate 100K samples from this distribution. Parallel MALA using quasi- with damping (called quasi- MLK in Gonzalez et al. [8]) converged to the ground truth sequential trace in 50 parallel iterations. The resulting set of samples at parallel iteration 50 are shown Figure 9, While here we visualized the results for quasi ELK, we note that quasi-DEER with clipped diagonal Jacobians between -1 and 1 also efficiently converged 72 parallel iterations. This preliminary result demonstrates the approaches can be used to traverse multimodal distributions.
28

---

Density
100K Parallel Samples (Iter 50)
Last 10K Samples
pa
50Q
coordinate 2
2
coordinate
eoordinate
80300 92000
94000
90000
58000
100200
coordinate
coordinate
iteration
Figure 9: Parallel sampling of multimodal mixture of Gaussians distribution. The left two panels show the target multimodal density and the final samples from running 50 parallel MALA iterations to simulate 100K samples. The rightmost panel shows the last 10K samples from sequential MALA sampling and parallel MALA sampling. Parallel MALA has converged to ground truth sequential trace of 100K samples in 50 parallel iterations such that the parallel and sequential traces are completely overlapping.
C Additional Results Figures
C.1 Parallel MALA
In this section, we include the following additional results figures for parallel MALA. The intent of these figures is to show trends over full parameter sweeps of batch size B and sequence length L that corroborate and augment our findings in the main text:
Figure 10: Wall-clock times (in seconds) for parallel MALA vs. sequential MALA sampling to generate batches of T, samples apiece. This contains the full results corresponding to Panel (A) in Figure 3 of the main text.
Figure 11: Distributions of iterations needed for each of B chains to numerically converge in parallel MALA. This contains the full results corresponding Panel (B) in Figure 3 of the main text.
Figure 12: MMD vs. wall-clock time performance profiles for parallel MALA vs. sequential MALA sampling. This contains the full results corresponding to Figure 4 of the main text.
Figure 13: MMD wall-clock time performance profiles for parallel MALA vs. sequential MALA sampling, interpolating over batchsize B. This figure provides an alternate perspec- tive (grouping by sequence length L instead of batch size B) of the results in Figure 12 above and Figure 4 of the main text.
Figure 14: Early-stopping MMD vs. wall-clock time tradeoffs. This figure contains the full results corresponding to Panel (B) in Figure 6 of the main text.
Figure 15: Maximum absolute errors accrued by parallel MALA w.r.t. the sequential sampling MALA algorithm function of Newton iteration. The purpose of this figure is to demonstrate that parallel MALA achieves numerical convergence on most combinations of batch size B and sequence length within the maximum number of iterations.
Figure 16: Pooled distributions of iterations needed for parallel MALA chains converge. This figure further emphasizes that the vast majority of length- L chains achieve numerical convergence well before exhausting the maximum number of allowed iterations.
Example samples of Bayesian logistic regression coefficient 31 on numerically- converged parallel MALA chain. This figure demonstrates that at numerical-comvergenet parallel and sequential MALA indeed output the same samples.
Figure 17:
Figure 18: Example samples of Bayesian logistic regression coefficient 31 on a not muuerically-converge parallel MALA chain. The purpose of this figure is demonstrate that even without complete numerical convergence, parallel MALA's 's generated samples are still effectively identical to those of sequential MALA's 's.
29

---

Wall-Clock Time Until Floating-Poine Convergence vs. Sequence Length (L)
==1
B
4
1.5
Time
0.5
0.0
=
8 16
B 32
16.0
Timee (sec)
5.0
0
0.0
10a
10a Sequence Length
30" Sequence Length
Length
Parallel MALA
Sequential MALA
90% MALA Interval 120x Seeds)
90% Seq. MALA Interval (20x Seeds)
Figure 10: Wall-clock times (in seconds) for parallel MALA sequential MALA sampling to generate B batches of L samples apiece. The lightly-colored bands represent 90% confidence intervals 20x random seeds. For all results shown above, parallel MALA was until full numerical convergence. This figure contains the full results corresponding to Panel (A) in Figure 3 of the main text. The main takeaway is that for all sequence lengths L and all batch sizes B except for B = 32, parallel MALA significantly faster than sequential MALA at generating samples.
Distributions of Numbers of Iterations Until Convergence Chain (t, for E (1, B})
8 1
8 2
B
Iters. Convg.
Pei Chain
60
60
60
40
40
20
20
10a
10'
302
10
10a
10
lters, to Convg.
(ta)
80
800
40
20
20
20
10a
10a
107
10a
107
10a
Sequence Length (L)
Sequence Length (L)
Sequence Length (L)
Example of Iters, ls (seed=13)
Median of team (20x Seeds)
60% CI of boss (20x Seeds! -- Max. Permitted Iters.
Figure 11: Distributions of iterations needed for each of B chains to numerically converge in parallel MALA. The red dotted line represents the maximum number of iterations permitted for given setting. For given collection of B chains, let represent the number of iterations needed for chain b of B to converge. Letting ((1) < ... < ((B) represent these numbers of iterations in increasing order, the solid blue line represents the median of f(0.9/) (i.e., the 90th percentile within these B chains) taken across 20x random seeds. The light-blue band represents 60% confidence interval of ((0.9B) computed across the 20x random seeds. The gray dots represent the observed numbers iterations until convergence for each chain in B for one representative seeded trial. This figure contains the full results corresponding to Panel (B) in Figure 3 of the main text. The main takeaway is that in most settings, parallel MALA converges in only a few dozen iterations despite generating tens, if not hundreds, of thousands of samples. Moreover, in the vast majority of settings, parallel MALA numerically converges before reaching the maæ_ é ter limit.
30

---

[Full-Convergence] MMD vs. Wall-Clock Time Tradeoff Curves + Interpolations (Fixing Batch Size B)
4
MMD (Lower is Better)
10"1
d
10-1
JR4
82K
Kak 426
-04
MHD (Lower is Better)
10-1
24
10.
10th
10
10t1
10
10P
Time (sec)
Time (sec)
Time (sec)
Parallel MALA
Interpolated Parallel MALA
Sequential MALA
Interpolated Sequential MALA
Figure 12: MMD wall-clock time performance profiles for parallel MALA sequential MALA sampling. The purpose of MMD is to measure sample quality. The main takeaway is that for all batch sizes except for B 32, we see that parallel MALA achieves virtually-ideniical MMD (i.e. outputs equally high quality samples) as sequential sampling, oftentimes an order of magnitude more faster. On B 32, achieve rough parity. The solid blue triangles and solid red circles represent the observed performances of parallel MALA and sequential MALA sampling, respectively, with corresponding sequence lengths L annotated next to each point. each panel, we also investigate the performance tradeoff's of using smaller sequence lengths, but over multiplc serial runs. For example, assuming parallel MALA run until numerical convergence, running parallel MALA with B 4 and L = 32K should output equivalent sample size and sample quality as running eight instances parallel MALA with = 4 and L 4K one after another, using the last samples from each serial run as the initialization for the next instance. Of course running eight instances of settings. However, using smaller sequence lengths could lead to proportionately faster numerical convergence per instance, making potentially more time-efficient certain cases than running one single instance with a longer sequence length. As such, we linearly interpolate to investigate such performance tradeoffs, with the interpolated values shown the hollow triangles and circles. For all panels, points closer to the bottom-left corner indicate better compute-time efficiency (i.e, faster sampling and higher quality of samples), This figure contains the full results corresponding to Figure 4 of the main text, organized by batch size B.
and L = 4K should take eight times longer than running one instance of the same
31

---

[Full-Convergence] vs. Wall-Clock Time Tradeoff + Interpolations (Fixing Sequence Length L)
1K
2K
L 4K
8K
MMD (Lower is Oetterl
10
10"4
10-3
10
10
10
10A
10
10
10
MMD (Lower better)
Parallel
10"4
Ex30
Interpolated Por. MALA
10-4
Sequential MALA
69.10
4x100
Interpolated
(sec)
Figure 13: MMD vs. wall-clock time performance profiles for parallel MALA vs. sequential MALA sampling, interpolating over batch size B. Why don't we always use B 32 (or the largest batch size possible)? Holding sequence length fixed each subplot, we explore MMD vs. wall-clock-time tradeoffs of using various batch sizes. For example, assuming parallel MALA is run until numerical convergence, running sixteen independent parallel MALA instances of and L 32K one after another should yield equivalent sample size and sample quality as running one parallel MALA instance of B = 32 and L = 32K, and take 16x longer than running one parallel MALA instance of B and L 32.. However, directly running B 32 and L 32K could overwhelm GPU memory, leading to significantly slower wall-clock performance than serially rumning smaller batches. This reasoning justifies our interpolation-over-BB analyses shown in the hollow triangles and circles. The interpolations in the L = 32K subplot reveal that running sixteen independent instances of parallel MALA at T. 32K iterations apiece would have achieved equivalent MMD as running parallel MALA or sequential sampling directly at B = 32 and 32K, but with noticeably faster speed. For all panels, points closer to the bottom-left corner indicate better compute-time efficiency (i.e, faster sampling and higher quality of samples). This figure provides an alternate perspective (grouping by sequence length L instead of batch size B) of the results in Figure 12 above and Figure 4 of the main text.
32

---

[Early-Stopping] MMD vs. Wall-Clock Time Tradeoff Curves Interpolations (Connecting by iteration)
2
MMD (Lower (ettter)
Par.
Far. (3x
30:3
shat.ge
(8x Por, (Full iters)
Afall 48th gic
Por (Pul itest
Far (Full Iters)
Sequertial
Sequarta
25
Sequential
10
10
10
10
10"
10P
10
10a2
10
10a
10
10n
10
101
8
8 16
MMD (Lower is Better)
(8x tars)
Itars
Par. (5x
10
-B Por. []1x Iters)
Por (10x tors)
-B- Far. (30x ftrs
-
Per: (16s (ters)
-0.
Per (20c hers)
-O-
Par. (20x ltens)
10
Por: (Full iters)
Por (Ful Ittr)
Far. (turt
Sequential
Sequentiin
Sequuntiial
al #/
Jogt oge 28.
dh- 0x_ g0r"
10"0
10ah
10a
10"0
10"
11P
10'
10ad
10"
(sec)
(sec)
(sec)
Figure 14: Early-stopping MMD vs. wall-clock time tradeoffs. Empirically, we know that the MMD of batch of parallel MALA sample chains converge that sequential MALA's in significantly fewer iterations than are required for the numerical convergence of individual chains. Each color-coded point represents the median MMD and interpolated wall-clock time (across 20x random seeds) of parallel MALA for a given batch size B and sequence length L after running for only few iterations (see individual subplot legends) vs. running until full numerical convergence ("Full Iters") ") and the sequential sampling baseline. The main takeaway is that in many cases, by performing early stopping (i.e., only running for very few parallel MALA iterations instead of waiting until full numerical convergence) one can obtain MCMC samples with comparable MMD sample quality but much faster time. This figure contains the full results corresponding Panel (B) in Figure 6 of the main text. For all panels, points closer to the bottom-left corner indicate better compute- time efficiency (i.e, faster sampling and higher quality samples).
33

---

Distribution of Maximum Absolute Errors Over Parallel MALA Iteration
8
=
8
8
16
=
-1k
10
10
30
70
10
1.0%
10
10
10'1
10"1-
20
L=2K
10
10"2
10
10
20
40
0
20
40
20
80
20
40
70
40
20
40
107
107
10-1
10.4
10"7
10-2
10-2
10-5
10-3
10-1
10-1
20
40
10
10t1
L=
10
10"7
10
10
10
10"1
10
L= 16K
10-4
10"4
10
10
10-1.
30
10"h
10
10-4
10
25
50
25
25
50
50
25
50
25
50
Iteration
1=32K
10--
10
30 0-
10
10
25
50
25
50
25
50
25
50
25
50
Iteration
-84K
1.0
109-
10
10a
10a1
10
10"1
10
72
50
Iteration
Iteration
Iteration
Iteration
Figure 15: Maximum absolute errors accrued by parallel MALA w.r.f. the sequential sampling MALA algorithm function of Newton iteration. The solid blue line represents the median of the median maximum absolute error (with the maximum taken over the entire sequence length L and dimensions D), with the inner median computed over all chains in B, and the outer median computed over 20x random seeds. The lower and upper limits of the light blue shaded region represent the medians (across 20x random seeds) of the 20th and 80th quantiles of the maximum absolute error (taken over the B chains per seed). The main takeaway of this figure is that parallel MALA achieves numerical convergence on most combinations of batch size B and sequence length within the maximum number iterations.
34

---

Distribution of Iterations Needed for Parailel MALA Converge
L= 1K
== 2K
-= 4K
1. BK
0.15
Frequency
0.30
0.08
64K
Frequency
0.30
0.09
0.00
Convg.
Iters. Needed Convg.
Needed Convg.
Figure 16: Pooled distributions of iterations needed for parallel MALA chains to converge. Each histogram shows the number of iterations required for convergence for all chains with sequence length L, pooled across all batch sizes and 20 random seeds. The red dotted line marks the 75th percentile, and the max_iter for each setting of L is marked by the right-most z-axis value. This figure emphasizes that the vast majority of length- I, chains achieve numerical convergence well before exhausting the maximum number of allowed iterations.
-04
-05
samples
p( x.
-0.9
5000
10000
15000
20000
25000
30000
-0.9
-0.8
-0.7
0066
-05
-04
Iteration
Parailel MALA
Sequential MALA
numerically- converged parallel MALA chain. Left: trace plots of 3, across iterations for sequential and parallel MALA sampling. The traces complete overlap, confirming that indeed have numerical convergence. Right: the histograms representing the posterior distributions implied by parallel and sequential MALA's samples identical.
Figure 17: Example samples of Bayesian logistic regression coefficient
35

---

-04
-0.6
saampees
p(8: X.
-0.8
-1.0
-1.2
5000
10000
15000
20000
25000
30000
-1.2
-1.0
-0.0
-0.4
Iteration
B1
Parallel MALA
Sequential MALA
Figure 18: Example samples of Bayesian logistic regression coefficient 81 not numerically- converged parallel MALA chain. Specifically, this chain reached max_ iters without numerical convergence. Left: trace plots of (3 across iterations for sequential and parallel MALA sampling. The traces still effectively completely overlap. implying that the two sets of samples are basically identical. The discrepancy convergence, the histograms representing the posterior distributions implied by parallel and sequential MALA's samples effectively identical. The main takeaway of this figure that even without complete numerical convergence, parallel MALA's 's generated samples are still effectively identical to those of sequential MALA's.
be seen around iteration 18000. Right: despite the lack of numerical
C.2 Parallel Leapfrog Integration for HMC
Sequential HMC Item-Response Time Comparison
Chains: 1, Step Size: 0.01
Chains: 2, Step Size: 0.01
1.75
time (8)
1.25
time
1.0
1.00
0.75
0.5
0.50
40
60
100
120
40
60
80
100
120
Number of Leapfrog Steps
Number of Leapfrog Steps
Chains: 1, Step Size: 0.02
Chains: 2, Step Size: 0.02
Sequenial Leapfrog
Parsiel Lespfrog
time P
time
1.5
1.0
0.5
40
80
100
120
40
100
Number of Leapfrog Steps
Number of Leapfrog Steps
Figure 19: Time comparison for sequential HMC the item-response model with sequential or parallel leapfrog integration. The time to simulate 1000 samples sequentially using HMC with either sequential parallel leapfrog the item-response model, function of the step size, number of leapfrog steps, and number of chains (panels). The target dimensionality is 1) = 501 for this problem and each likelihood evaluation is computed across approximately 30K datapoints. For one chain, HMC with parallel leapfrog is fastest across the two step sizes. However at two chains we do see speed-ups ups when using sequential HMC with parallel leapfrog for this problem.
36


Reviews:
[{"review_id": "Cl7YcLVPjL", "reviewer": "Reviewer_Eern", "summary": "Previous work on parallelising MCMC on multi-core CPUs and GPUs has focused on running multiple chains in parallel. This method takes recent work on parallel evaluation of non-linear recursions using Newton's method (DEER, quasi-DEER) and proposes a framework to run an MCMC sampler in parallel across the chain length. The authors demonstrate their approach by extending Gibbs sampling, Metropolis-adjusted Langevin, and Hamiltonian Monte Carlo. They then extend their framework to use quasi-Newton methods which reduces the wall-clock runtimes by orders of magnitude for a number of test problems.", "questions": "1) Do you see a clear path towards extending your framework to support adaptive MCMC samplers? If so, could you include a sentence or two about what would be required? It is not clear whether this is feasible.", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths\n\n1) Very explanation of the problem, and demonstration of their solutions effectiveness with a number of widely-used MCMC samplers.\n2) Significantly reduces wall-clock time for a number of MCMC samplers by permitting the use of greater computational resources.\n3) Good evaluation of the method for a number of test problems, including Bayesian Logistic Regression.\n\nWeaknesses\n\n1) Method cannot clearly be applied directly to adaptive MCMC samplers.\n2) Speedup is less than linear relative to hardware, so 2x hardware doesn't lead to 2x performance improvement. Although this is mitigated through the availability of GPUs.", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n**Adaptive MCMC Samplers (Question 1, Weakness 1)**\n\nExtending 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.\n\nNUTS 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). \n\nOn 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. \n\nAside 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. \n\n**Relationship between wall-clock speedup and hardware (Weakness 2)**\n\nHere 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. \n\n**Demonstration of approach in higher-dimensions**\n\nFinally, 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."}, {"review_id": "bi266WQ3R6", "reviewer": "Reviewer_G15U", "summary": "MCMC has a computational cost that is linear to the number of samples that one would like to simulate. This computational cost can be improved via parallelization. Most existing methods does so by simulating several independent Markov chains and aggregate all resulting samples. This paper takes a different approach and simulates a sequence of MCMC states (given a fixed length) by treating the sequence as the solution to a fixed-point problem, which can be solved using a parallelized form of Newton's method. The author specifies the nonlinear sequence model for a number of MCMC methods, which can be used to obtain the MCMC state sequence as the solution of corresponding fixed-point problems. It is shown that this approach can lead to improved wall-clock times compared to the traditional sequential simulation of MCMC states when the number of states generated is the same.", "questions": "One suggestion that I have is to expand the background section. While the author does provide all necessary citations to understand the DEER algorithm (which is a core component of the proposed methodology), the text in section 2 itself may not be detailed enough for those that have not heard of the DEER algorithm before (myself included). \n\nMore specifically, it wasn't immediately clear to me that the number (T) of MCMC states to be simulated is fixed, and that the Newton iterations are optimizing the entire sequence of T states as a whole. It may be a good idea to outline the DEER algorithm in section 2. This would help make the subsequent sections of the paper easier to read.", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 4, "originality": 3, "strengths_and_weaknesses": "This paper is for the most part very well organized and clearly written. The problem studied here is very important as improving the computational cost of MCMC is a long-standing topic. The proposed algorithm here is easy to understand and simple to use. Compared to other methods that combine states from multiple independent MCMC instances, there is no extra care needed for aggregating samples from different chains.", "was_revised": true, "final_justification": "This paper looks at parallelizing MCMC from a different but interesting angle, and the authors addressed my questions during the rebuttal period. I am keeping my positive score.", "rebuttal": "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.\n\n**Additional Background on DEER**\n\nWe 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. \n\n**Extending approach to adaptive MCMC Samplers**\n\nAn 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. \n\n**Demonstration of approach in higher-dimensions** \n\nFinally, 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."}, {"review_id": "FTh8CsiGTa", "reviewer": "Reviewer_m4iF", "summary": "Previous work has shown that nonlinear RNNs can be evaluated in parallel using Newton's\nmethod (DEER) and this has been made stable and faster in subsequent work (ELK).\nThe current paper applies this technique to the task of drawing MCMC samples thus\nenabling withing-chain parallelism.\n\nThe paper also explores the application of various standard approximations to further\nbring down the computation cost to a product of the state space times the length of the\nsequence (as opposed to quadratic in the size of the state space).\n\nFurther the paper demonstrates that the application of DEER (and its approximations) to\nMCMC chains leads to very fast convergence. Thus we don't need a very long sequence of\nstates.", "questions": "On line 94 what does it mean to say \"g <- \\sigma(g) + ... \"? (This looks like a recurrence?)\n\nHow big is the posterior state space in the German Credit Bayesian Logistic Regression \nexample?\n\nHow does this method work on problems with multimodal posteriors? Does the small number of\nsamples that seem to be needed for DEER also allow for efficient exploration of \nmultimodal posteriors.", "limitations": "yes", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths - novel application of RNN literature to MCMC\n\nThe original DEER and ELK were designred for parallelizing nonlinear RNNs. The authors\nhave taken that work and drawn an interesting parallel to MCMC. This appears to be\nnovel and allows for MCMC to leverage advances in a different area of AI.\n\n## Strengths - fast convergence and high ESS.\n\nThe experiments do seem to suggest that the application of DEER to MCMC allows for \nconvergence with very few samples and these tend to have very high Effective Sample Size\n(ESS). This makes the approach very promising.\n\n## Weakness - unclear MALA algorithm\n\nI couldn't follow the use of the stop_gradient in the MALA algorithm. I do understand\nhow this trick normally works, but this formulation of the MALA algorithm was not very\nclear to me and Appendix A doesn't have any further explanation even though line 101 \npromises that appendix A would have additional details.\n\n## Weakness - very small scale experiments\n\nThe N schools example only had an 18-dim posterior. Given the significant improvements\nclaimed in the paper the authors could have demonstrated on a higher dimenstional problem.\n\n## Weakness - no demonstration of accuracy\n\nThe paper has proposed a number of approximations but none of the experiments demonstrate\nthat the results are actually correct.\n\nFurther none of the models have multimodal posteriors so it is unclear whether all of these\napproximations will allow for efficient exploration of the posterior space.", "was_revised": true, "final_justification": "The authors claim of improving over the performance of variational methods on unimodal non-gaussian posteriors is not demonstrated to a convincing degree in this paper. MCMC methods need to demonstrate excellence in multimodal posteriors which this paper doesn't attempt to demonstrate. The authors tease about some results in the rebuttal which cannot be validated. In any case, if the authors were aware of other MCMC papers with specific multimodal posteriors used for evaluations then those should have been front and center in their evaluations.", "rebuttal": "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. \n\n**MALA algorithm details** \n\nWe 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.\n\n**Accuracy of algorithms**\n\nThe 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.\n\nFirst, 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. \n\nSecond, 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.\n\nWe 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. \n\n**Dimensionality of experiments and extending to higher dimensions**\n\nWe 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.  \n\nSpecifically, 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. \n\nWe 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.\n\n**Multimodal posteriors**\n\nThe 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. \n\nHere 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.   \n\n**Conclusion**\n\nWe 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."}, {"review_id": "gKMRLK8Fop", "reviewer": "Reviewer_hixK", "summary": "This paper tackles the problem of speeding up Markov chain Monte Carlo (MCMC) by parallelizing across the sequential steps of each chain. The authors adapt the DEER algorithm (a parallel Newton-based method originally used for RNN-style recurrences) to solve an entire chain’s updates in parallel. This approach is novel (up to my best knowledge): by performing $O(\\log T)$ parallel Newton iterations instead of $T$ sequential steps, they show it’s possible to generate long trajectories with only a few dozen iterations. The key innovation is recasting common MCMC samplers (Gibbs, MALA, HMC) as fixed-point problems that can be solved using parallel Newton-type methods.\n\nThe authors develop several algorithms within the DEER framework:\n- **DEER**: Uses full Jacobians for quadratic convergence with O(log T) parallel time complexity\n- **Quasi-DEER**: A cheap variant using only Jacobian diagonals \n- **Stochastic Quasi-DEER**:  reduces the memory footprint by estimating Jacobian diagonals\n- **Block Quasi-DEER**: Optimized for HMC's leapfrog integrator structure\n\nThe methods demonstrate  10×–30× speedups over sequential sampling across various applications, including hierarchical Gaussian models, Bayesian logistic regression, and complex targets. \n\nThis breakthrough transforms MCMC from a fundamentally sequential O(T) process into a parallelizable O(log T) problem, opening new possibilities for scaling Bayesian inference on modern accelerators. The work establishes both theoretical foundations and practical algorithms for leveraging massive parallelism in Monte Carlo sampling.", "questions": "-  The paper references the DEER algorithm (from Lim et al. [5]) as part of the parallelization approach, but its presentation is quite limited. Even if DEER is not the central theme, a clearer self-contained explanation of what DEER does and how it integrates into your method would greatly improve readability. For instance, please elaborate on the intuition and steps of the DEER algorithm (and quasi-DEER) in the main text or appendix. This helps readers who are unfamiliar with prior work [5,6] understand the foundation of your parallel MCMC approach without extensive external reading. \n- The paper makes several claims (e.g. about convergence guarantees, sublinear scaling, “orders of magnitude” speedups, etc.) that would benefit from stronger theoretical backing or at least references to supporting theory. I encourage the authors to support these claims with background results or analysis. For example, if you claim that the stochastic quasi-DEER “enjoys provable convergence guarantees” or that Newton’s method yields quadratic convergence, please either cite the precise theorems/propositions (from prior work or your appendix) or outline the conditions under which these claims hold. \n- While I find the application of parallel Newton/DEER methods to MCMC to be novel and exciting, the manuscript could do more to explicitly position this contribution relative to prior work. I suggest the authors clearly highlight what is new in this work (e.g. parallelizing specific MCMC samplers across chain length, the new quasi-Newton variants, etc.) versus what is adopted from existing methods. For instance, discuss how your approach differs from or improves upon prior parallel MCMC efforts or parallel diffusion simulations [18–22,28,29] mentioned in related work.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- The paper builds on a existing framework (treating the chain update as a fixed-point problem solved by Newton’s method) and extends it with appropriate modifications. For example, the authors acknowledge that the straightforward DEER algorithm would require storing and multiplying through $T$ Jacobian matrices of size $D\\times D$, leading to cubic time and quadratic memory in state dimension $D$. They address this by using quasi-Newton updates that only use the diagonal of each Jacobian (or Hutchinson’s stochastic diagonal estimate), reducing complexity to $O(D)$ and making the method much more scalable. They also introduce a block-diagonal Newton scheme for HMC’s leapfrog integrator, which accounts for position–momentum coupling and roughly halves the number of iterations needed to converge compared to a naive diagonal approach.\n-  the authors identified potential bottlenecks (memory use, slow convergence) and devised solutions grounded in numerical methods. The approach for handling non-differentiable Metropolis acceptance steps is also reasonable: they perform the exact accept/reject in the forward pass but omit its gradient in the backward pass (stop-gradient), citing a proposition that guarantees Newton convergence despite this discontinuity. While this is a somewhat heuristic workaround, it’s backed by theory and appears to work well in practice (the parallel MALA yields correct results). Overall, the methodology is sound and the authors seem aware of its theoretical conditions and limitations\n- The authors implement their parallel MCMC algorithms in JAX on  GPU, ensuring a comparison with optimized baselines. They test the approach on multiple problems: a hierarchical Bayesian model (the 18-dimensional “eight schools” problem) using a reparameterized Gibbs sampler, a Bayesian logistic regression (BLR) with MALA and HMC on the German Credit dataset, and even a 501-dimensional item-response model (in the appendix) to demonstrate scaling to higher dimensions. For each case, the paper reports  raw speed, sample quality via Maximum Mean Discrepancy (MMD) to a ground-truth or long-run baseline. The numerical results strongly support the paper’s claims. For instance, in the BLR experiment, parallel MALA converges in only tens of Newton iterations even for very long chains. The Gibbs sampler results show a more modest but still solid gain – roughly 2× speedup over traditional sampling on the eight-schools model when chains are long. The benefit there is limited (likely by the low dimension and cheaper updates per step), but it demonstrates that even a relatively simple Gibbs chain can leverage GPU parallelism for moderate gains. \n\nWeaknesses\n- The main weaknesses of the paper lie in the limitations and scope of the proposed approach. First, while the method does achieve large speedups, it comes at the cost of substantially increased computation per iteration and memory usage. The authors explicitly note that their algorithms “trade memory and compute for time”. In scenarios where the target density or gradient is extremely expensive to evaluate, performing many parallel gradient evaluations (as this method does) might not actually be efficient – the paper acknowledges that if the model is very costly per step, those resources might be better spent on more sequential steps or independent chains.\n- Secondly, the memory overhead is non-trivial. Even with the quasi-Newton diagonal approximation, parallelizing very long chains or using large batch sizes can exhaust GPU memory (the authors report out-of-memory errors for their largest tested setting, 32 parallel chains of length 64K, using the memory-hungry version). The paper’s solution – the memory-efficient Jacobian estimator – was critical to make the BLR experiments feasible, but this comes at the price of using stochastic estimates, which might in some cases slow convergence (indeed, they observed the plain quasi-Newton often took the maximum allowed iterations to converge for HMC until the block method was introduced). Thus, while the authors did mitigate this issue, scalability to very high dimensions or extremely long chains remains a concern \n- One minor concern is that the analysis of convergence and conditions for quadratic speed-up are mostly cited from prior work; the paper itself does not deeply analyze when Newton’s method will converge quickly for a given MCMC (aside from empirical observation). However, given space constraints, this is understandable and the empirical results provide evidence that a few dozen Newton iterations usually suffice. \n- The paper focuses on Gibbs (with continuous reparameterizable conditionals), Langevin dynamics, and HMC with fixed hyperparameters. It does not address more complex adaptive methods like NUTS or discrete-variable samplers. The authors explicitly leave extension to adaptive HMC and other MCMC variants for future work. Consequently, the immediate applicability is somewhat limited – many practical Bayesian inference tasks use NUTS or other adaptive schemes that would need non-trivial adaptation of this method. Finally, from an originality perspective, one could argue that the paper is an incremental advance in that it leverages an existing technique (parallel Newton solver from prior work) rather than inventing a wholly new sampling paradigm. The true novelty is in recognizing that this technique can be applied to MCMC and making it work for those algorithms; this is a meaningful innovation, but it’s not a completely new theory of MCMC.", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n**Elaboration on DEER algorithm (Question 1)**\n\nWe 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. \n\n**Theoretical Claims and Guarantees (Question 2, Weakness 3)**\n\nOur 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.\n\nFor 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.\n\n**Clarity of contributions and Novelty (Question 3, Weakness 4)**\n\nWe 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. \n\n**Scalability (Weakness 2)**\n\nThe 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. \n\nFirst, 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.\n\nNext, 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. \n\n**Significance of Advance and Adaptive MCMC (Weakness 4)**\n\nWe 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. \n\nHowever, 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."}]

Author rebuttal:

Meta-review: Standard MCMC is inherently sequential, with cost scaling linearly in the number of samples, which makes it poorly suited to parallel hardware. Prior work mostly runs many chains in parallel (ie. embarrassingly parallel setting), leaving the sequential bottleneck untouched. This paper asks whether a single chain can be parallelized across its length, turning an O(T) process into something closer to O(log T). This direction could make MCMC practical at scales where sequential simulation is too slow.

The main contribution is to show that Gibbs, MALA, and HMC can be written as fixed-point recurrences and evaluated in parallel using Newton-type methods, building on the DEER (non-linear Differential Equation as fixed point itERation) framework. The authors propose practical variants such as diagonal and stochastic Jacobian approximations and block updates for HMC, which reduce memory and computation and make the method usable on GPUs. They also show how the Metropolis step can be handled by exact evaluation in the forward pass with gradients omitted in the backward pass. Experiments demonstrate significant speedups, including on higher-dimensional problems. The novelty is in adapting DEER-style methods to MCMC and making them practical for standard samplers.

During rebuttal, reviewers asked for more details, scaling to higher dimensions, treatment of adaptive samplers, the accept–reject step, and performance on multimodal targets. The authors added a clearer description of the method, fixed a presentation error in MALA, and reported new experiment. These additions convinced two reviewers and softened concerns of a third, though one remained unconvinced about multimodal performance and theoretical depth.

Overall, there is agreement that the idea is novel and tackles a long-standing limitation of MCMC. While theory remain incomplete, the authors strengthened the paper and showed feasibility on challenging problems. The work opens a promising research direction.
Decision: Accept (poster)