Running Markov Chain Monte Carlo on Modern Hardware and Software

TL;DR

This paper explores how modern hardware and software, like GPUs and JAX, can be leveraged to significantly accelerate Markov chain Monte Carlo (MCMC) methods, which have lagged behind deep learning in hardware adoption.

Contribution

It introduces practical patterns and insights for adapting MCMC algorithms to modern parallel hardware and software, demonstrating substantial speedups over traditional CPU-based approaches.

Findings

Significant speedups achieved with GPU and TPU implementations.

Effective use of deep learning frameworks like PyTorch and JAX for MCMC.

Discussion of potential pitfalls in hardware acceleration of MCMC.

Abstract

Today, cheap numerical hardware offers huge amounts of parallel computing power, much of which is used for the task of fitting neural networks to data. Adoption of this hardware to accelerate statistical Markov chain Monte Carlo (MCMC) applications has been much slower. In this chapter, we suggest some patterns for speeding up MCMC workloads using the hardware (e.g., GPUs, TPUs) and software (e.g., PyTorch, JAX) that have driven progress in deep learning over the last fifteen years or so. We offer some intuitions for why these new systems are so well suited to MCMC, and show some examples (with code) where we use them to achieve dramatic speedups over a CPU-based workflow. Finally, we discuss some potential pitfalls to watch out for.