用有限状态机优化并行马尔可夫链蒙特卡洛，提升计算效率。
Efficiently Vectorized MCMC on Modern Accelerators
- 将单链MCMC改写为有限状态机，避免并行时的同步等待
- 实测加速比最高达10倍，理论推导支持最优设计
- 适合需要高效采样的统计建模与贝叶斯推断用户
随着自动向量化工具（如JAX的vmap）的出现，编写多链马尔可夫链蒙特卡洛（MCMC）算法已变得极为简便，只需对单链代码应用这些工具即可。然而，对于多种MCMC算法，这种方法会导致同步开销——在每轮迭代中，所有并行运行的链必须等待最慢的链完成采样。本文提出通过有限状态机（FSM）框架重构单链MCMC算法，使向量化时无需同步等待。基于简化模型，推导出该方法可获得的理论加速比，并据此给出算法设计的合理建议。我们实现了椭圆切片采样、HMC-NUTS和延迟拒绝等主流算法的FSM版本，在实验中实现最高达一个数量级的加速。
原文摘要 · Abstract (English)
With the advent of automatic vectorization tools (e.g., JAX's $\texttt{vmap}$), writing multi-chain MCMC algorithms is often now as simple as invoking those tools on single-chain code. Whilst convenient, for various MCMC algorithms this results in a synchronization problem -- loosely speaking, at each iteration all chains running in parallel must wait until the last chain has finished drawing its sample. In this work, we show how to design single-chain MCMC algorithms in a way that avoids synchronization overheads when vectorizing with tools like $\texttt{vmap}$ by using the framework of finite state machines (FSMs). Using a simplified model, we derive an exact theoretical form of the obtainable speed-ups using our approach, and use it to make principled recommendations for optimal algorithm design. We implement several popular MCMC algorithms as FSMs, including Elliptical Slice Sampling, HMC-NUTS, and Delayed Rejection, demonstrating speed-ups of up to an order of magnitude in experiments.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。