Parallelizing MCMC Across the Sequence Length

TL;DR

This paper introduces parallel algorithms for MCMC sampling across sequence length, enabling significant speedups by solving fixed-point problems with parallel Newton methods, thus overcoming the sequential nature of traditional MCMC.

Contribution

The authors develop novel parallel algorithms for MCMC that evaluate samplers across sequence length using fixed-point formulations and parallel Newton methods, improving efficiency.

Findings

Achieved up to hundreds of thousands of samples with tens of Newton iterations.

Parallel algorithms accelerated MCMC sampling by over an order of magnitude in several examples.

Developed two new parallel quasi-Newton methods with lower memory and runtime.

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.