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

David Yallup; Namu Kroupa; Will Handley

Source-linked AI summary

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

David Yallup, Namu Kroupa, Will Handley

arXiv:2601.23252v2Jan 30, 2026stat.COcs.LGstat.ML

TL;DR

Evidence estimation and posterior sampling remain difficult to scale for multimodal targets because Nested Sampling is sequential and its hard-constrained kernels have irregular computation. The paper introduces GPU-oriented Nested Slice Sampling with Hit-and-Run Slice Sampling and slice-width tuning, and reports robust performance across challenging inference tasks, especially multimodal ones.

Problem
Nested Sampling is robust for complex targets but its sequential structure and data-dependent constrained updates map poorly to accelerator hardware.
Method
NSS vectorizes the Nested Sampling outer update, uses Hit-and-Run Slice Sampling for constrained propagation, and tunes slice width near a derived optimum to stabilize per-step costs.
Results
Across challenging synthetic targets, high-dimensional Bayesian inference, and GP hyperparameter marginalization, NSS shows competitive performance and robust evidence estimation with high-quality posterior samples.
Takeaways & Limitations
NSS is a practical default for multimodal or constrained-support targets when evidence estimation and posterior samples are both required, while gradient-based methods remain preferable for smooth unimodal targets.
Takeaways & Limitations
NSS can be slower in wall-clock time when strict SIMD synchronization penalizes variable-cost inner updates, although near-optimal slice widths reduce this effect.

Abstract

from arXiv · show

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.

1 Introduction

The paper targets evidence estimation and posterior inference for complex distributions by making Nested Sampling efficient on GPU accelerators. It introduces a vectorized implementation using Hit-and-Run Slice Sampling, principled slice-width tuning, and open-source software.

Evidence estimation supports Bayesian model comparison, while standard MCMC methods generally do not estimate the marginal likelihood directly.
Existing Nested Sampling formulations are difficult to accelerate because classical updates are sequential and constrained kernels have data-dependent control flow.
The implementation vectorizes the full Nested Sampling update, including energy evaluation, thresholding, resampling, and mutation for GPU hardware.
Hit-and-Run Slice Sampling provides the constrained inner kernel, while near-optimal slice-width tuning concentrates per-step costs and improves batching on accelerators.
Experiments compare NSS with adaptive tempered SMC baselines on multimodal synthetic targets, high-dimensional Bayesian inference, and machine-learning inference tasks.
NSS and its underlying Nested Sampling framework are released as composable open-source software for probabilistic programming ecosystems.

2 Background and Theoretical Framework

Nested Sampling estimates a normalizing constant by progressively restricting a reference density to lower-energy regions and applying quadrature over shrinking prior volumes. Its practical bottleneck is constrained sampling within feasible regions that can become small, anisotropic, or disconnected, motivating slice-based updates.

2.1 Nested Sampling: Nested Sampling transforms normalizing-constant estimation into one-dimensional quadrature over automatically selected energy levels and associated prior volumes.
2.1 Nested Sampling: The algorithm iteratively removes the highest-energy live particles and replaces them with samples from the reference density restricted below the current energy threshold.
2.1 Nested Sampling: Batch removals can be unrolled into sequential single-death events, with each removal contributing its own order-statistic volume compression.
2.2 Constrained sampling: As constraints tighten, feasible regions may become small, anisotropic, and disconnected, making the constrained sampler central to practical efficiency and correctness.
2.2 Constrained sampling: Region-based rejection sampling is effective in low-to-moderate dimensions, whereas constrained MCMC avoids global proposal envelopes but can suffer boundary rejections and dimension-dependent difficulties.
2.2 Constrained sampling: Hit-and-Run Slice Sampling selects a random direction and performs an exact one-dimensional slice update along a chord of the constrained set.
2.3 Connections to particle methods: Nested Sampling resembles particle methods that bridge reference and target distributions, but it estimates normalization through probabilistic prior-volume quadrature rather than standard SMC importance weights.

3 Algorithm and implementation

NSS vectorizes Nested Sampling by batching deletion, resampling, constrained mutation, and replacement, with HRSS providing the constrained updates. Its tuning analysis and GPU experiments support predictable costs, strong scaling, and practical batch-size and slice-width rules, while identifying geometric assumptions and scope limits.

Algorithm and implementation: NSS separates a vectorized outer kernel from a pluggable constrained-update kernel targeting the truncated reference distribution.The outer kernel maintains live particles and evidence bookkeeping, while replacement kernels generate constrained samples.
Algorithm and implementation: Each batched iteration deletes the k worst points, duplicates surviving parents, mutates duplicates under the energy constraint, and replaces the deleted points.The batch ratio k/m controls expected compression, while k exposes parallelism independently of live-set resolution m.
Nested Slice Sampling: NSS uses HRSS for constrained propagation, combining random directions with exact one-dimensional slice updates along constrained chords.This kernel is selected for high-dimensional constrained regions that may be compact, anisotropic, sharply bounded, or disconnected.
Nested Slice Sampling: The fully vectorized outer kernel maps batched constrained mutation and likelihood evaluation to GPU execution, including synchronized parallel replacement chains.The implementation is presented as a GPU-oriented form of Nested Sampling rather than primarily CPU-core parallelism.
Optimal slice width tuning: The tuning analysis gives w∗≈1.36ℓ for fixed chord length and high-dimensional optimal width scaling as d−1/2, with anisotropy entering mainly through the mean eigenvalue.At the optimum, per-step cost variability remains modest, supporting vectorization of many short constrained chains.

4 Experiments

Across synthetic, Bayesian, and GP benchmarks, NSS produces competitive posterior and evidence estimates, with especially strong performance on multimodal and hierarchical targets. The experiments compare NSS with adaptive tempered SMC variants and posterior-only baselines using distributional, evidence, and predictive metrics.

Benchmark design: NSS and adaptive tempered SMC are evaluated on synthetic multimodal and hierarchical targets, Inference Gym models, and GP hyperparameter marginalization.The comparisons include SMC-RW, SMC-IRMH, SMC-SS, SMC-HMC, and, where appropriate, SS and NUTS.
Synthetic benchmarks: On the 2d 40-mode Gaussian mixture, all particle methods recover the correct mode structure, while larger performance gaps appear on harder targets.The harder cases involve high-dimensional multimodality, invalid likelihood regions, and hierarchical structure.
Synthetic benchmarks: NSS achieves the lowest MMD on the 10d funnel and generally provides the most reliable mode recovery on the 10d multimodal Gaussian mixture.Slice sampling embedded in particle methods, particularly NSS, is reported as reliable for mode recovery.
Inference Gym: On Inference Gym problems, NSS matches strong baselines on lower-dimensional models, remains competitive on higher-dimensional models, and achieves the best performance across all metrics on RadonIndiana.All NSS posteriors are broadly consistent with NUTS reference draws; SMC-RW substantially underestimates evidence on RadonIndiana.
GP hyperparameter marginalization: NSS provides the best predictive performance on Mauna Loa and competitive performance on Airline, while maintaining high effective sample size and similar predictive trends across methods.The GP evaluation uses held-out test NLL and RMSE, with 95% credible intervals shown for NSS posterior predictive distributions.

5 Discussion

The discussion identifies tuning and hardware execution as central practical considerations for NSS. It also frames the method as useful for multimodal or constrained targets while noting limits in wall-clock performance and very high dimensions.

Practical tuning: The chain length p trades off parent-particle correlation against wasted computation from over-mixing.This trade-off applies to mutation steps in NSS and the SMC baselines.
Practical tuning: Near the optimal slice width, HRSS step costs concentrate, making batched execution more effective than naive constrained random walks.This reduces the impact of variable-cost stepping-out and shrinkage updates under lockstep execution.
Hardware trade-offs: NSS can be slower in wall-clock time than SMC on hardware where strict SIMD synchronization penalizes variable-cost inner updates.The method may still be efficient in energy evaluations per effective sample despite this hardware-dependent cost.
Scope: The discussion presents NSS as a natural default for multimodal or constrained-support targets, simulator-like or non-differentiable likelihoods, and settings requiring evidence alongside posterior samples.This conclusion is stated within the scope of the chosen benchmarks and the method’s constrained-sampling design.
Limitations: For d > 10^3, more sophisticated constrained mutation kernels may be required, while exact evidence estimation for highly multimodal non-convex targets remains challenging.The paper identifies practical scaling limits as an open problem.

6 Conclusion

The paper concludes with an accelerator-oriented Nested Sampling implementation that combines a batched outer kernel with Hit-and-Run Slice Sampling. NSS is tuned for vectorized execution and demonstrates competitive performance across challenging inference settings, with open-source code and several remaining research directions.

Contribution: NSS instantiates a batched Nested Sampling outer kernel with Hit-and-Run Slice Sampling for hard constraints and massively parallel likelihood evaluation.The implementation is designed for accelerator-oriented inference.
Methodological contribution: The slice-width analysis gives the near-optimal rule w*≈1.36 ℓ and links this regime to well-behaved per-step costs for vectorized execution.The analysis also includes high-dimensional scaling for ellipsoidal sets.
Empirical scope: NSS shows competitive performance on challenging synthetic targets, high-dimensional Bayesian inference, and GP hyperparameter marginalization.These applications use evidence estimation for model comparison and uncertainty-related evaluation.
Outlook: The open-source implementation aims to make Nested Sampling easier to use and benchmark in ML settings while remaining compatible with established practice.Open problems include automatic mutation-length selection, asynchronous accelerator execution, and gradient-informed constrained updates.

Code Availability

The implementation is available as a current working repository release, while the paper benchmarks adaptive tempered SMC using ESS-controlled temperature schedules and standard mutation kernels.

A current working Nested Sampling implementation and experiment scripts are available in the authors’ GitHub repository.
Tempered SMC benchmarks use bridging distributions with β_0 = 0 and β_T = 1, incremental weights, resampling, and MCMC mutation.
The next temperature is selected adaptively so that the effective sample size is approximately ρm, and log Z accumulates from log-mean estimates.
The baselines use mutation kernels and tuning rules matched closely to NSS, including comparable particle counts and inner-kernel steps.

A.2 Adaptive non-equilibrium samplers and parallelism

The paper frames nested and SMC samplers as adaptive non-equilibrium algorithms whose data-dependent execution complicates accelerator use, then specifies their parallel outer and inner kernels and comparison settings.

Adaptive execution: Adaptive nested and SMC samplers evolve through distributions with data-dependent termination and non-deterministic outer-loop lengths, making accelerator execution less immediate.
Adaptive execution: The implementation evaluates adaptive-loop overhead against fixed scans and AIS on two 20-dimensional problems, including host–accelerator communication and schedule discovery.
Results: On the Gaussian target, all methods recover similar evidence, whereas SMC-IRMH and AIS substantially underestimate evidence on the mixture target; NSS is more expensive but more robust there.
Nested Sampling kernel: Nested Sampling evaluates energies, deletes the highest-energy particles in batches, resamples surviving parents, and mutates replacements under a hard energy constraint.
NSS configuration: Default NSS uses m = 1000, k = 100, and p = d, with slice width scaled through adaptive whitening of the live-set covariance.
SMC baselines: SMC baselines use adaptive tempering, high ESS targets, and random-walk, independent-MH, HMC, or slice-sampling mutation kernels.
Scope: The comparison excludes learned proposals and parallel tempering to focus on strong, standard kernels, while reporting point estimates of log Z for SMC.

A.6 Hit-and-Run Slice Sampling Implementation

NSS uses Hit-and-Run Slice Sampling to make constrained updates, combining exact one-dimensional slice moves with bounded loops and adaptive geometry handling for predictable accelerator execution.

HRSS update: HRSS updates a state along a randomized line, treating points outside the hard constraint or prior support as having log density −∞.
Stepping-out: Stepping-out expands a randomized bracket until it encloses the horizontal slice, with a ten-expansion safety cap for JIT-compatible bounded execution.
Shrinkage: Shrinkage repeatedly samples within the bracket and contracts it when needed, producing an exact idealized slice update without a separate Metropolis–Hastings step.
Safeguards: The 100-iteration shrinkage cap can return a null move under pathological conditions, but neither cap was reached in the section 4 experiments.
Precision: HRSS relies only on density comparisons, enabling single or mixed precision when log-density computation remains stable; performance was robust in the reported experiments.
Parallel execution: Concentrated per-step cost matters for GPU batches because the slowest parallel chain determines wall-clock time and otherwise leaves faster chains waiting.
Cost concentration: NSS maintains approximately 5 evaluations per constrained step with low variance across dimensions, while rejection-based random walks have higher and more variable cost.
Cost concentration: Near the optimal slice width, HRSS cost variance remains O(1) across a broad dimensional range because stepping-out adapts to local geometry.

B Extended Experimental Results

The extended experiments evaluate NSS on multimodal, hierarchical, Bayesian, and Gaussian-process problems, finding accurate evidence estimates, strong mode recovery, and substantial accelerator scaling.

Evaluation setup: Synthetic evaluations report W2 and MMD quality metrics, density evaluations, wall-clock time, and ESS per wall-clock time across independent runs and seeds.
Synthetic benchmarks: The 40-component Gaussian mixture has analytically known log evidence ln Z = −9.21 and is used to assess multimodal sampling quality.
Synthetic benchmarks: The 10-dimensional five-component Gaussian mixture tests mode discovery and weighting across separated regions, with ground-truth samples drawn directly from its components.
Funnel benchmark: The 10-dimensional funnel couples x-scale exponentially to y, while all tested algorithms use the difficult centered parameterization against non-centered HMC ground truth.
Inference Gym: With proper initialization, NSS has evidence variance comparable to SMC-HMC and substantially outperforms SMC-RW on the S&P500 stochastic-volatility model.
Gaussian-process inference: The Gaussian-process experiment uses a composite constant, linear, and spectral-mixture kernel with 11 hyperparameters and constrained ordered frequencies.
Evidence accuracy: NSS evidence errors remain approximately 0.2–0.35 across dimensions, while combined estimates reach 0.12 at 10d and 0.02 at 20d.
Mutation scaling: NSS plateaus near p ≈7 with MMD ≈0.11, motivating the consistent choice p = d while noting that the optimal chain length is problem-specific.

C.1 Comparison with existing Nested Sampling implementations

NSS is compared with CPU-oriented Nested Sampling implementations on the Mauna Loa Gaussian-process hyperparameter problem using matched hardware, population, termination, and likelihood settings. It achieves substantially faster execution while also improving predictive and evidence-related outcomes.

Architectural comparison: Existing implementations such as PolyChord and UltraNest target CPU architectures and rely on MPI-based parallelism, complicating fair comparison with GPU-accelerated NSS.Their additional features are described as poorly suited to GPU execution.
Experimental setup: The comparison uses the same population size, termination criterion, NVIDIA A100 hardware, and JIT-compiled likelihood across NSS, UltraNest, and JAXNS.UltraNest uses SliceSampler and vectorized PopSliceSampler configurations closest to NSS.
Results: 12× to 30× speedup over UltraNest is achieved by NSS depending on configuration, alongside substantially better predictive NLL and fewer likelihood evaluations.The comparison concerns Mauna Loa Gaussian-process hyperparameter marginalization.
Results: Over 10× slower than NSS, JAXNS also produces substantially worse evidence estimates and predictive metrics despite operating in the same JAX ecosystem.The likelihood is relatively bottlenecked by Cholesky factorization, so larger speedups are expected for more parallelizable likelihoods.

D.1 Sample Quality Metrics

The paper evaluates posterior sample quality with distributional discrepancies and predictive metrics, while describing batched NSS updates and geometric uncertainty propagation. These diagnostics distinguish sample agreement, predictive calibration, and point-prediction accuracy, but geometric uncertainty excludes imperfect constrained mixing.

Distributional metrics: MMD measures the distance between two sample sets’ mean embeddings, equaling zero exactly when characteristic-kernel samples come from the same distribution.The implementation uses a Gaussian RBF kernel with bandwidth set by the median heuristic.
Distributional metrics: Sliced 2-Wasserstein distance projects samples onto random one-dimensional directions and computes exact one-dimensional Wasserstein distances along those projections.The method is computationally efficient and statistically robust in higher dimensions, using 200 random projection directions.
Predictive metrics: NLL evaluates how well posterior predictive distributions explain held-out observations, whereas RMSE evaluates point-prediction accuracy.NLL and RMSE together distinguish methods with similar point estimates but different uncertainty quantification.
Batched updates: Each batched NSS iteration removes k highest-energy live points and generates k replacements under the shared k-th-worst-energy threshold.Unrolling deletions into k single-death events gives the associated joint prior-volume contraction under the usual i.i.d. constrained-prior idealisation.
Uncertainty propagation: Geometric uncertainty is propagated by Monte Carlo shrinkage simulation, but this uncertainty estimate excludes additional errors from imperfect constrained mixing.Repeated runs with different seeds and initial conditions complement the geometric estimate in selected experiments.
Tuning analysis: The HRSS tuning analysis derives expected computational cost, a unique fixed-slice optimum, and asymptotic slice-width scaling relevant to NSS.These results target predictable per-step cost for constrained sampling.

E.1.2 Convergence of expectations in a ball

This section analyzes scaled random variables derived from a uniformly sampled point in a high-dimensional ball and establishes their limiting behavior through Gamma-distribution representations. It then uses independence, moment bounds, and uniform integrability to control convergence of expectations and related quadratic-form quantities.

Ball decomposition: A uniformly sampled point in the d-dimensional unit ball is decomposed into an independent radius and uniformly distributed direction.The analysis defines directional and radial components used to construct the scaled variables.
Joint convergence: Independence of the finite-dimensional scaled components is preserved under joint convergence through factorization of characteristic functions.The resulting sum converges by the continuous mapping theorem to a Gamma-distributed limit.
Expectation convergence: Uniform integrability is established for the relevant scaled variables, allowing convergence in distribution to support convergence of expectations.The argument uses moment bounds and comparisons between the constructed variables.
Quadratic-form control: For positive-definite quadratic forms with spectra uniformly bounded away from zero and infinity, variance and Taylor-remainder bounds are uniform in dimension.The resulting error control depends only on spectral bounds rather than dimension.
Ellipsoid connection: The analysis applies these bounds to normalized chord lengths of ellipsoids, whose high-dimensional behavior is governed by the quadratic form induced by the ellipsoid matrix.The ellipsoid chord-length construction transforms the problem to the unit ball.

E.1.4 Optimal scaling of hit-and-run slice sampling in ellipsoids

The section derives the computational cost of Hit-and-Run Slice Sampling and optimizes the slice width for fixed intervals and high-dimensional ellipsoids. The resulting rule scales inversely with the square root of dimension, while numerical tests agree with the theory.

Cost model: Hit-and-Run Slice Sampling restricts each direction to the connected slice component containing the starting point, an interval of length ℓ.The expected cost is measured by likelihood evaluations per sampling step.
Cost model: Theorem 1 gives conditional expectations for stepping-out and shrinkage steps, yielding the expected number of likelihood evaluations per sampling step.The shrinkage analysis models repeated uniform proposals within the current bracket until the slice is hit.
Fixed-slice optimization: Theorem 2 establishes a unique global minimizer for the expected cost on a fixed slice, with optimal width proportional to the slice length.The optimum is written as w∗(ℓ) = u∗ℓ.
Ellipsoid scaling: κ∞≈1.3035 determines the asymptotic optimal-width constant for high-dimensional ellipsoids.The dimension-dependent constant satisfies κd = κ∞ + o(1).
Ellipsoid scaling: The optimal width in an ellipsoid scales as O(d−1/2), reflecting concentration of probability mass in a high-dimensional ball.The leading geometry enters through the ellipsoid’s mean chord length.
Numerical validation: Numerical tests agree with the theory within statistical error, while standard deviation remains order 1 across dimensions up to d = 1000.For the cube, standard deviation increases with dimension, though slowly.

F.2 Numerical comparison

The numerical comparison evaluates constrained samplers across increasing dimensionality and multimodality. Slice Sampling remains accurate in high dimensions, while reflection-based alternatives develop systematic evidence errors as dimension increases.

The test varies α ∈[0, 1], where increasing α increases multimodality and sampling difficulty.α = 0 represents a unimodal Gaussian, whereas α = 1 represents a highly multimodal likelihood.
Figure 13 compares constrained samplers in dimensions d ∈{2, 4, 10, 30, 100, 300} using 1000 live points.
Slice Sampling remains correct across all tested dimensions, including the high-dimensional settings.
GMC and GMC-2019 systematically deviate above 30 dimensions, while RSS produces incorrect evidence already in four dimensions.
The reflection-based samplers yield systematically negative evidence because their constrained updates create boundary overdensities rather than uniform within-contour densities.GMC-2019 deviates less strongly because its trajectories do not leave the boundary and therefore shift density toward the centre of the level set.