[Paper Review] Pigeons.jl: Distributed Sampling From Intractable Distributions

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

Gaussian Processes and Bayesian Inference6 citations

TL;DR

Pigeons.jl is a Julia package that enables distributed, reproducible sampling from complex, intractable probability distributions—such as multimodal posteriors or high-dimensional distributions—using a novel approach to strong parallelism invariance. It guarantees identical outputs across single-threaded, multi-threaded, and distributed (MPI) executions for the same seed, ensuring scientific reproducibility and simplifying debugging of randomized, parallel MCMC algorithms.

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.

Motivation & Objective

• To address the challenge of reproducibility in distributed and parallel MCMC sampling, where output typically varies with the number of threads or processes.
• To provide a scalable, user-friendly interface for sampling from complex, high-dimensional, or multimodal distributions common in Bayesian inference and statistical physics.
• To enable accurate estimation of both expectations ∫f(x)π(x)dx and normalization constants Z, which are often inaccessible in existing tools.
• To support arbitrary data types, including combinatorial objects like phylogenetic trees, beyond standard continuous or discrete parameter spaces.
• To ensure correctness and ease of validation through strong parallelism invariance (SPI), making debugging and testing deterministic across execution configurations.

Proposed method

• The package implements distributed, non-reversible parallel tempering as the core sampling algorithm, enabling efficient exploration of complex target distributions.
• It uses the SplittableRandoms.jl package to generate independent, reproducible random number streams from a single seed, ensuring determinism across thread and process counts.
• The implementation leverages Julia’s multi-threading and MPI interfaces to support execution on single machines, multi-core systems, and clusters with thousands of nodes.
• Users specify the unnormalized density γ(x), and Pigeons.jl returns samples for computing expectations or estimating the normalization constant Z.
• Strong parallelism invariance is enforced by deterministically splitting a master random number generator across all threads and processes, eliminating non-deterministic behavior in random number generation.
• The software provides a simple, high-level API that abstracts away low-level distributed computing complexity while preserving full control over execution configuration.

Experimental results

Research questions

• RQ1 Can a distributed MCMC sampler guarantee identical output across different numbers of threads and processes for the same random seed?
• RQ2 How can strong parallelism invariance be achieved in a high-performance, distributed, and randomized algorithm without sacrificing scalability?
• RQ3 What are the key sources of non-determinism in parallel and distributed random number generation that undermine reproducibility in scientific software?
• RQ4 Can a general-purpose sampling package efficiently handle both continuous and combinatorial parameter spaces while supporting normalization constant estimation?
• RQ5 How does the integration of splittable random number generation improve the correctness and testability of distributed MCMC algorithms?

Key findings

• Pigeons.jl guarantees strong parallelism invariance: for a given seed, the output is identical regardless of the number of threads or processes used, enabling deterministic debugging and validation.
• The package successfully supports distributed execution across thousands of MPI-communicating machines, with identical results across single-threaded, multi-threaded, and distributed configurations.
• By using SplittableRandoms.jl, Pigeons.jl achieves independent, reproducible random number streams across all execution units, eliminating non-determinism in random number generation.
• The framework enables accurate estimation of both posterior expectations and normalization constants Z, which are often omitted or approximated in existing tools.
• The implementation demonstrates that high-performance, scalable, and reproducible sampling is achievable in practice, even for high-dimensional and multimodal posteriors.
• The package is compatible with Julia 1.8+ and integrates with multiple probabilistic programming languages, enabling broad applicability across scientific domains.