[论文解读] Pigeons.jl: Distributed Sampling From Intractable Distributions

Nikola Surjanovic, Miguel Biron-Lattes+4 more|arXiv (Cornell University)|Aug 18, 2023

Gaussian Processes and Bayesian Inference被引用 6

一句话总结

Pigeons.jl 是一个 Julia 软件包，通过一种新颖的强并行性不变性方法，实现了对复杂、不可解概率分布（如多模态后验或高维分布）的分布式、可复现采样。它保证了在相同种子下，单线程、多线程和分布式（MPI）执行的输出完全一致，从而确保了科学可复现性，并简化了随机化并行 MCMC 算法的调试。

ABSTRACT

We introduce a software package, Pigeons.jl, that provides a way to leverage distributed computation to obtain samples from complicated probability distributions, such as multimodal posteriors arising in Bayesian inference and high-dimensional distributions in statistical mechanics. Pigeons.jl provides simple APIs to perform such computations single-threaded, multi-threaded, and/or distributed over thousands of MPI-communicating machines. In addition, Pigeons.jl guarantees a property that we call strong parallelism invariance: the output for a given seed is identical irrespective of the number of threads and processes, which is crucial for scientific reproducibility and software validation. We describe the key features of Pigeons.jl and the approach taken to implement a distributed and randomized algorithm that satisfies strong parallelism invariance.

研究动机与目标

• 解决分布式和并行 MCMC 采样中的可复现性挑战，因为输出通常随线程数或进程数而变化。
• 为贝叶斯推断和统计物理中常见的复杂、高维或多重模态分布，提供可扩展且用户友好的采样接口。
• 实现对期望值 ∫f(x)π(x)dx 和归一化常数 Z 的准确估计，这些在现有工具中通常不可访问。
• 支持任意数据类型，包括组合对象（如系统发育树），而不仅限于标准的连续或离散参数空间。
• 通过强并行性不变性（SPI）确保正确性与可验证性，使不同执行配置下的调试和测试具有确定性。

提出的方法

• 该软件包以分布式、不可逆的并行退火作为核心采样算法，能够高效探索复杂的目标分布。
• 它使用 SplittableRandoms.jl 包，从单一种子生成独立且可复现的随机数流，确保在不同线程和进程数量下保持确定性。
• 实现利用 Julia 的多线程和 MPI 接口，支持在单机、多核系统以及拥有数千个节点的集群上执行。
• 用户指定非归一化密度 γ(x)，Pigeons.jl 返回用于计算期望值或估计归一化常数 Z 的样本。
• 通过确定性地将主随机数生成器拆分到所有线程和进程中，强制实现强并行性不变性，从而消除随机数生成中的非确定性行为。
• 该软件提供简单、高层级的 API，抽象了底层分布式计算的复杂性，同时保留了对执行配置的完全控制。

实验结果

研究问题

• RQ1 一个分布式 MCMC 采样器能否保证在相同随机种子下，无论使用多少线程或进程，输出完全一致？
• RQ2 如何在高性能、分布式且随机化的算法中实现强并行性不变性，同时不牺牲可扩展性？
• RQ3 在并行和分布式随机数生成中，导致科学软件可复现性被破坏的主要非确定性来源是什么？
• RQ4 一个通用采样包能否高效处理连续和组合参数空间，同时支持归一化常数估计？
• RQ5 可分割随机数生成的集成如何提升分布式 MCMC 算法的正确性和可测试性？

主要发现

• Pigeons.jl 保证了强并行性不变性：对于给定种子，无论使用多少线程或进程，输出完全一致，从而实现确定性调试与验证。
• 该软件包成功支持在数千台通过 MPI 通信的机器上进行分布式执行，且在单线程、多线程和分布式配置下结果完全相同。
• 通过使用 SplittableRandoms.jl，Pigeons.jl 在所有执行单元中实现了独立且可复现的随机数流，消除了随机数生成中的非确定性。
• 该框架能够准确估计后验期望值和归一化常数 Z，而这些在现有工具中通常被忽略或近似处理。
• 该实现表明，即使在高维和多模态后验的情况下，高性能、可扩展且可复现的采样在实践中是可行的。
• 该软件包兼容 Julia 1.8 及以上版本，并与多种概率编程语言集成，可在科学领域的广泛应用中发挥效用。