Nested Slice Sampling: Vectorized Nested Sampling for GPU-Accelerated Inference

TL;DR

The paper tackles scalable Bayesian inference and model comparison for complex, multimodal targets by reimagining Nested Sampling for GPUs. It introduces Nested Slice Sampling (NSS), a vectorized NS where the inner constrained updates use Hit-and-Run Slice Sampling, yielding stable, parallelizable performance. A principled optimal slice-width tuning and online whitening stabilize per-step costs, enabling effective batching on accelerators. Empirical results show NSS achieves accurate evidence estimates and high-quality posterior samples across synthetic, Inference Gym, and Gaussian Process marginalization tasks, often outperforming tempered SMC baselines and existing CPU-oriented NS codes. The work provides an open-source NSS implementation that facilitates adoption and reproducibility in ML and scientific computing contexts.

Abstract

Model comparison and calibrated uncertainty quantification often require integrating over parameters, but scalable inference can be challenging for complex, multimodal targets. Nested Sampling is a robust alternative to standard MCMC, yet its typically sequential structure and hard constraints make efficient accelerator implementations difficult. This paper introduces Nested Slice Sampling (NSS), a GPU-friendly, vectorized formulation of Nested Sampling that uses Hit-and-Run Slice Sampling for constrained updates. A tuning analysis yields a simple near-optimal rule for setting the slice width, improving high-dimensional behavior and making per-step compute more predictable for parallel execution. Experiments on challenging synthetic targets, high dimensional Bayesian inference, and Gaussian process hyperparameter marginalization show that NSS maintains accurate evidence estimates and high-quality posterior samples, and is particularly robust on difficult multimodal problems where current state-of-the-art methods such as tempered SMC baselines can struggle. An open-source implementation is released to facilitate adoption and reproducibility.

Figures (12)

• Figure 1: Nested Sampling shrinks a reference density by successive energy constraints (left), yielding a quadrature approximation of the normalizing constant (right).
• Figure 2: Posterior samples for a mixture of 40 bivariate Gaussians on a bounded uniform prior. NSS (orange) recovers all modes with correct weights. This 2$d$ example, despite its low dimensionality, challenges many sampling methods due to the large number of well-separated modes.
• Figure 3: Posterior predictive distributions for GP regression on Mauna Loa CO$_2$ [left] and Airline passengers [right]. Shaded regions show 95% credible intervals from NSS samples, demonstrating well-calibrated uncertainty quantification.
• Figure 4: Distribution of likelihood evaluations per constrained step for NSS (slice sampling, blue) vs RW (random walk, orange) on an ill-conditioned Gaussian ($\kappa=100$) in dimensions 10, 50, and 100. NSS maintains concentrated cost ($\approx 5$ evals, with standard deviation $\approx 1$) while RW exhibits heavy tails that degrade batched execution efficiency.
• Figure 5: Marginal distribution of the first two parameters of the $d{=}10$ mixture of Gaussians. Ground truth samples (blue) compared with NSS samples (orange).

Theorems & Definitions (35)

• Lemma 1
• proof
• Lemma 2
• proof
• Lemma 3
• proof
• Lemma 4
• proof
• Lemma 5
• proof