Parallelizing MCMC Across the Sequence Length
David Zoltowski, Skyler Wu, Xavier Gonzalez, Leo Kozachkov, Scott Linderman
Markov chain Monte Carlo (MCMC) methods are foundational algorithms for Bayesian inference and probabilistic modeling. However, most MCMC algorithms are inherently sequential and incur computational costs that scale linearly with sample size. 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.* 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 the widely-used Gibbs sampling, Metropolis-adjusted Langevin, and Hamiltonian Monte Carlo algorithms can be evaluated in parallel using this approach. We showcase several examples in which we can generate hundreds of thousands of MCMC samples with only a few dozen parallel Newton iterations. Additionally, we develop two new quasi-Newton methods for the evaluation of nonlinear recursions that lower memory costs and reduce runtime. Across multiple challenging benchmarks, we find that the proposed parallel algorithms can accelerate MCMC chains by orders of magnitude compared to sequential evaluation.