📊 statistics

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

This paper introduces Nested Slice Sampling (NSS), a GPU-accelerated, vectorized formulation of Nested Sampling that leverages Hit-and-Run Slice Sampling and an optimized slice width rule to efficiently handle complex, multimodal Bayesian inference problems while maintaining accurate evidence estimates and posterior samples.

Original authors: David Yallup, Namu Kroupa, Will Handley

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: David Yallup, Namu Kroupa, Will Handley

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the best possible treasure map for a complex, mountainous island. You don't know exactly where the treasure is, but you have a rough idea of the terrain (the "prior") and a way to measure how good a spot looks (the "likelihood"). Your goal is twofold:

Map the terrain: Find all the places where the treasure might be hidden (the "posterior").
Count the islands: Figure out exactly how many distinct treasure islands exist and how big they are (the "evidence" or "marginal likelihood"). This helps you decide if you should even bother looking for treasure on this island or switch to a different one.

This paper introduces a new, super-fast way to do this mapping and counting, specifically designed to run on modern, powerful computer chips called GPUs (the same kind used to train AI and play video games).

Here is the breakdown of their method, Nested Slice Sampling (NSS), using simple analogies:

1. The Problem: The "Sequential" Bottleneck

Traditional methods for this task (like Nested Sampling) work a bit like a single-lane road.

You have a group of explorers (particles) looking for treasure.
To find the best spots, you have to send them out one by one, check their location, and then send the next one.
Because they have to wait for each other, this process is slow, especially when the terrain is tricky (multimodal) or very high-dimensional (like a map with thousands of directions).
Furthermore, the rules for moving around are "hard constraints" (you can't go into a swamp). Traditional methods struggle to move efficiently inside these tight, weirdly shaped swamps.

2. The Solution: The "Highway" Approach

The authors built Nested Slice Sampling (NSS), which turns that single-lane road into a massive, multi-lane highway.

Vectorization: Instead of sending explorers out one by one, NSS sends out thousands of them at the exact same time.
GPU Power: It uses the parallel processing power of GPUs to evaluate the "goodness" of all these spots simultaneously.
The Result: It's like going from a snail's pace to a rocket ship. The paper shows it can be 10 to 30 times faster than existing methods on the same hardware.

3. The Secret Sauce: "Hit-and-Run" Slice Sampling

How do you move thousands of explorers efficiently inside a tight, weirdly shaped swamp without them getting stuck?

The Old Way: Some methods try to guess a path, but if they hit a wall, they bounce back and try again. In high dimensions, this leads to a lot of wasted time (bouncing off walls).
The NSS Way (Hit-and-Run): Imagine an explorer standing in the swamp. They pick a random direction (like throwing a dart) and draw a straight line through the swamp. They then pick a new spot uniformly along that line.
Why it works: This method is incredibly good at navigating complex shapes. It doesn't get stuck in corners or bounce endlessly. It finds a valid new spot quickly, no matter how twisted the terrain is.

4. The "Goldilocks" Tuning Rule

One of the biggest headaches in these methods is choosing the "step size" (how far the explorer jumps).

Too small: You take tiny steps and never get anywhere.
Too big: You jump over the valid area and waste time recalculating.
The Paper's Discovery: The authors did the math to find the perfect step size. They found a simple rule: The ideal step size depends on the "width" of the valid area.
The Analogy: It's like tuning a radio. If you are too close to the station, the signal is weak; too far, and it's static. They found the exact frequency where the signal is clearest.
Why it matters for speed: When you use this perfect step size, every explorer takes roughly the same amount of time to finish their step. This is crucial for GPUs. If one explorer takes 100 steps and another takes 1, the GPU has to wait for the slow one. By making everyone's step time consistent, the whole highway moves at maximum speed.

5. What They Tested

The authors didn't just talk about theory; they tested it on:

Synthetic Traps: They created fake, extremely difficult landscapes with many separate "islands" of treasure (multimodal problems) and deep, narrow funnels. NSS found all the islands and mapped them correctly, while other methods got stuck or missed spots.
Real-World Data: They tested it on real datasets (like predicting airline passengers or credit risk) and found it produced accurate maps and counts, often beating other state-of-the-art methods.
Speed: On a standard high-end GPU, NSS was dramatically faster than older methods, proving that you can run these complex Bayesian calculations in seconds rather than hours.

Summary

Nested Slice Sampling is a new tool that takes a difficult, slow mathematical process (finding treasure in complex, high-dimensional spaces) and makes it fast, parallel, and robust.

It uses GPUs to do thousands of calculations at once.
It uses a clever "Hit-and-Run" strategy to navigate tricky terrain without getting stuck.
It uses a mathematical rule to ensure everyone moves at the same speed, preventing bottlenecks.

The result is a method that is particularly great for problems where the "map" is full of separate islands (multiple solutions) or has strange, hard-to-navigate shapes, making it a powerful new tool for scientists and machine learning engineers.

Technical Summary: Nested Slice Sampling (NSS)

Problem Statement

Bayesian model comparison and calibrated uncertainty quantification require integrating over parameter spaces to estimate the marginal likelihood (evidence), $Z$ . While Nested Sampling (NS) is a robust alternative to standard Markov Chain Monte Carlo (MCMC) for this task, particularly on complex, multimodal targets, its traditional implementation faces significant hurdles in modern computing environments. Classical NS is inherently sequential, relying on dynamic, data-dependent control flow (e.g., variable iteration counts for constrained sampling) that maps poorly to Single Instruction, Multiple Data (SIMD) hardware like GPUs. Existing implementations (e.g., PolyChord, MultiNest) often rely on CPU-centric features such as clustering heuristics, MPI-based parallelism, and unit-hypercube transforms, which are ill-suited for accelerator hardware. Furthermore, constrained inner kernels often suffer from high variance in computational cost per step, leading to inefficient batched execution where faster chains wait for slower ones.

Methodology

The paper introduces Nested Slice Sampling (NSS), a vectorized formulation of Nested Sampling designed for GPU acceleration. The methodology consists of two primary components: a parallel outer loop and a specific constrained inner kernel.

1. Vectorized Outer Loop

Unlike traditional NS which removes one live point at a time ( $k=1$ ), NSS employs a batched deletion strategy. In each outer iteration, the algorithm:

Evaluates energies for a population of $m$ live particles.
Identifies the $k$ worst particles (highest energy) and sets a batch threshold $E_{batch}$ .
Resamples $k$ parent particles from the surviving $m-k$ population.
Applies a constrained update kernel to these $k$ parents in parallel to generate new particles satisfying $E(x) < E_{batch}$ .
Replaces the dead particles with the new samples.

This structure exposes parallelism across the $k$ mutation chains, allowing the entire update (energy evaluation, thresholding, resampling, and mutation) to be compiled into a single JIT-compiled function targeting GPU hardware.

2. Hit-and-Run Slice Sampling (HRSS) as the Inner Kernel

To address the inefficiencies of constrained MCMC on accelerators, NSS utilizes Hit-and-Run Slice Sampling (HRSS) as the constrained update kernel.

Mechanism: HRSS selects a random direction and performs a one-dimensional slice update along the chord of the constrained set using "stepping-out" and "shrinkage" procedures.
Tuning Analysis: A key theoretical contribution is the derivation of an optimal slice width $w^*$ . The authors prove that for high-dimensional ellipsoidal level sets, the optimal width scales as $w^* \propto d^{-1/2}$ .
Cost Concentration: Crucially, the analysis shows that operating near this optimal width causes the computational cost (number of likelihood evaluations per step) to concentrate with low variance. This concentration mitigates SIMD divergence, making batched execution of HRSS chains highly efficient compared to rejection-based methods or random walks, which exhibit heavy-tailed cost distributions.

3. Implementation

The implementation is built within the JAX ecosystem using the Blackjax library. It treats the outer NS update and inner constrained update as composable functions, enabling automatic differentiation and efficient compilation to CPU, GPU, and TPU backends. The framework avoids legacy complexities like clustering heuristics, relying instead on the geometric compression of the NS path and the robustness of slice sampling.

Key Contributions

Massively Parallel Implementation: A fully vectorized Nested Sampling framework where the outer kernel and batched constrained mutations map directly to parallel GPU execution.
Principled Tuning of Slice Width: A theoretical analysis yielding a near-optimal rule for setting the slice width ( $w^* \approx 1.36 \ell$ for fixed chords, scaling as $d^{-1/2}$ for ellipsoids). This tuning ensures per-step costs concentrate, enabling effective vectorization.
Robustness on Difficult Targets: The use of HRSS within the NS framework provides a robust alternative to gradient-based methods on non-differentiable or highly multimodal targets where standard MCMC or SMC baselines struggle.
Open-Source Release: An open-source implementation designed to be composable within modern probabilistic programming ecosystems.

Experimental Results

The authors evaluated NSS against adaptive tempered Sequential Monte Carlo (SMC) baselines (using RW, IRMH, SS, and HMC kernels) and standalone NUTS/SS on three categories of problems:

Synthetic Benchmarks:
- Multimodality: On a 2D mixture of 40 Gaussians and a 10D mixture of 5 Gaussians, NSS successfully recovered all modes with correct weights, outperforming SMC-RW and SMC-HMC in high-dimensional multimodal settings.
- Hierarchical Structure: On Neal's 10D funnel, NSS achieved the lowest Maximum Mean Discrepancy (MMD) and Sliced 2-Wasserstein distance ( $W_2$ ), demonstrating superior mixing in hierarchical geometries.
Inference Gym (Real-Data Problems):
- On lower-dimensional problems (EightSchools, GermanCredit), NSS matched or exceeded the posterior quality of SMC baselines.
- On higher-dimensional problems (S&P500, RadonIndiana), NSS remained competitive. Notably, on the RadonIndiana hierarchical model, SMC-RW failed to estimate evidence correctly, while NSS provided tight, accurate estimates.
Gaussian Process Hyperparameter Marginalization:
- Applied to Mauna Loa CO2 and Airline Passenger datasets with a spectral mixture kernel ( $d=11$ ).
- NSS provided the best predictive performance (lowest Test NLL) on the Mauna Loa dataset and competitive results on Airline Passengers.
- Scaling: GPU scaling experiments demonstrated nearly linear speedup with batch size $k$ (up to $k/m \le 0.5$ ) and near-constant runtime scaling as population size $m$ increased, confirming efficient hardware utilization.
- Comparison: NSS achieved a 12x to 30x speedup over existing NS implementations (UltraNest, JAXNS) on identical hardware while maintaining or improving evidence accuracy and predictive metrics.

Significance and Claims

The paper claims that Nested Slice Sampling represents a practical, end-to-end GPU-viable implementation of Nested Sampling. Its significance lies in:

Bridging the Gap: It successfully translates the sequential, constraint-heavy logic of NS into a form compatible with modern accelerator hardware, a domain where NS has historically lagged behind SMC and MCMC.
Robustness: It offers a reliable default for problems with hard constraints, non-differentiable likelihoods, or complex multimodal geometries where gradient-based methods (like HMC) may fail or require extensive tuning.
Theoretical Insight: The derivation of the optimal slice width and the demonstration of cost concentration provide a principled basis for tuning constrained samplers in high dimensions, explaining why vectorized execution is feasible despite the adaptive nature of slice sampling.
Reproducibility: By providing an open-source, composable implementation, the authors aim to facilitate adoption in both physical sciences and machine learning, serving as a strong baseline for future comparisons with neural samplers.

The authors remain modest, acknowledging that for smooth, unimodal targets with reliable gradients, gradient-based alternatives (SMC-HMC, NUTS) may still be more efficient. They also note that scaling to very high dimensions ( $d > 10^3$ ) may eventually require more sophisticated, gradient-guided constrained samplers.