Parallelizing MCMC Across the Sequence Length
David M. Zoltowski, Skyler Wu, Xavier Gonzalez, Leo Kozachkov, Scott W. Linderman
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers1
Ask how each one uses itBuilds on12
- Your classifier is secretly an energy based model and you should treat it like oneWill Grathwohl, Kuan-Chieh Wang, Jörn-Henrik Jacobsen, David Duvenaud et al.ICLR 2020 · 643 citations
- ADAHESSIAN: An Adaptive Second Order Optimizer for Machine LearningZhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa et al.AAAI 2021 · 358 citations
- Parallel Sampling of Diffusion ModelsAndy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh et al.NeurIPS 2023 · 144 citations
- Simplified State Space Layers for Sequence ModelingJimmy T. H. Smith, Andrew Warrington, Scott W. LindermanICLR 2023 · 78 citations
- Variational Bayesian Last LayersJames Harrison, John Willes, Jasper SnoekICLR 2024 · 75 citations
Related papers
- Involutive MCMC: a Unifying FrameworkKirill Neklyudov, Max Welling, Evgenii Egorov, Dmitry P. VetrovICML 2020 · 40 citations
- Structured Stochastic Gradient MCMCAntonios Alexos, Alex J. Boyd, Stephan MandtICML 2022 · 14 citations
- Towards Scalable and Stable Parallelization of Nonlinear RNNsXavier Gonzalez, Andrew Warrington, Jimmy T. H. Smith, Scott W. LindermanNeurIPS 2024 · 47 citations
- Distributed Metropolis Sampler with Optimal ParallelismWeiming Feng, Thomas P. Hayes, Yitong YinSODA 2021 · 7 citations
- PROCA: Programmable Probabilistic Processing Unit Architecture with Accept/Reject Prediction & Multicore Pipelining for Causal InferenceYihan Fu, Anjunyi Fan, Wenshuo Yue, Hongxiao Zhao et al.HPCA 2025 · 3 citations
