Source-linked AI summary
Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro
Du Phan, Neeraj Pradhan, Martin Jankowiak
TL;DR
Probabilistic programming needs a NumPy-based interface that can retain Pyro’s modeling abstractions while exploiting JAX’s functional transformations. NumPyro composes effect handlers with JAX transformations and introduces iterative NUTS for end-to-end JIT compilation. The resulting implementation is much faster than existing alternatives across small and large dataset regimes.
Problem
NumPyro addresses how Pyro’s modeling API can operate on a fundamentally different JAX-based functional backend while retaining acceleration, differentiation, and vectorization capabilities.
Method
NumPyro combines Pyro-compatible effect handlers with JAX transformations and converts NUTS tree construction from recursive to iterative form for end-to-end compilation.
Results
NumPyro’s end-to-end JIT-compiled NUTS is much faster than existing implementations in both small and large dataset regimes.
Takeaways & Limitations
NumPyro provides Pyro’s modeling language while leveraging JAX transformations to accelerate and vectorize probabilistic inference subroutines.
Abstract
from arXiv · showhide
NumPyro is a lightweight library that provides an alternate NumPy backend to the Pyro probabilistic programming language with the same modeling interface, language primitives and effect handling abstractions. Effect handlers allow Pyro's modeling API to be extended to NumPyro despite its being built atop a fundamentally different JAX-based functional backend. In this work, we demonstrate the power of composing Pyro's effect handlers with the program transformations that enable hardware acceleration, automatic differentiation, and vectorization in JAX. In particular, NumPyro provides an iterative formulation of the No-U-Turn Sampler (NUTS) that can be end-to-end JIT compiled, yielding an implementation that is much faster than existing alternatives in both the small and large dataset regimes.
1 Introduction
NumPyro embeds probabilistic programming in Python and NumPy while using JAX for automatic differentiation, vectorization, and JIT compilation. Its effect handlers compose with JAX tracing so Pyro’s modeling interface can run on this backend.
- NumPyro is a probabilistic programming package built atop JAX that lets users write programs with familiar NumPy arrays and operations.
- JAX provides program transformations including automatic differentiation, vectorization, and JIT compilation for Python and NumPy functions.
- Effect handlers inject computation into primitive probabilistic-program statements, such as recording random choices in an execution trace.
- NumPyro composes effect handlers with JAX tracing and control-flow primitives while preserving Pyro’s modeling language.
2 Support for Pyro’s Modeling Interface
NumPyro preserves Pyro’s modeling and inference interface while adapting effect handling to a functionally different JAX backend. This shared abstraction supports common probabilistic-program transformations across the two systems.
- NumPyro retains Pyro’s language primitives and modeling and inference interface, including sample and param statements.
- NumPyro provides trace, replay, and condition effect handlers for nonstandard interpretations of sample and param statements.
- Effect handling supplies a composable common interface despite fundamental differences between PyTorch’s dynamic style and JAX’s functional tracing requirements.
3 Leveraging JAX Transformations in Inference Subroutines
NumPyro uses composable JAX transformations and effect handlers to accelerate inference subroutines. It converts NUTS tree construction to an iterative form for end-to-end JIT compilation and uses vmap for batching.
- JAX transformations: JAX offers grad, jit, and vmap transformations, which NumPyro applies to purely functional inference subroutines for optimized parallel execution.
- Iterative NUTS: The recursive NUTS BuildTree formulation is difficult to JIT compile because JAX cannot trace its complex recursive control flow.
- Iterative NUTS: NumPyro converts BuildTree into an iterative procedure so NUTS can be fully JIT compiled and accelerated through XLA.
- Iterative NUTS: An end-to-end compiled iterative NUTS implementation avoids repeated Python dispatch overhead across many LeapFrog steps and enables greater XLA operator fusion.
- Vectorizing with vmap: vmap represents batched computations as mapping over an outermost argument axis without requiring changes to the underlying code.
- Vectorizing with vmap: Composing vmap with seed, trace, and condition supports vectorized prior sampling, posterior prediction, log-likelihood computation, and SVI estimation.
4 Experiments
NumPyro’s end-to-end JIT-compiled NUTS was evaluated against Stan and Pyro across small and large-data models, with strong speedups reported in several settings.
- Benchmark setup: NumPyro’s NUTS was compared with Stan and Pyro on an HMM, logistic regression, and SKIM across small and large data regimes.NumPyro uses end-to-end JIT compilation, while Pyro compiles only potential-energy computation.
- Hidden Markov Model: 340X faster than Pyro and 6X faster than Stan, NumPyro achieved these HMM speedups in a CPU-only comparison.The small dataset was expected to perform poorly on GPU, so the comparison used CPUs only.
- Logistic Regression: About 2X faster than Pyro, NumPyro performed better on logistic regression with more than half a million datapoints.GPU acceleration significantly outperformed CPU execution, while large tensor operations narrowed differences among GPU implementations.
- Sparse Kernel Interaction Model: NumPyro had consistently lower time-per-effective-sample overhead than Stan as SKIM dimensionality increased.The comparison examined NUTS scaling with dataset dimensionality for a sparse regression model with pairwise interactions.
- Effective sample size: 788, 556, and 652 were the average effective sample sizes for NumPyro 64-bit, NumPyro 32-bit, and Stan, respectively.These averages came from five runs with different random seeds, 1000 warmup steps, and 1000 samples per run.
- Precision and hardware: Single precision had lower time per effective sample because double precision’s higher effective sample size did not offset its longer inference time.The result concerns the evaluated NumPyro precision settings.
5 Summary
NumPyro combines Pyro-compatible effect handling with JAX transformations for accelerated probabilistic inference. Its iterative NUTS implementation delivers strong CPU performance for small models and GPU performance for larger ones.
- Summary: NumPyro combines JAX transformations with effect handlers to provide Pyro’s modeling language while accelerating and vectorizing inference.Effect handlers remain transparent to the JAX tracer and composable with its transformations.
- Summary: An iterative NUTS formulation enables end-to-end JIT compilation and strong performance on CPUs for small models and GPUs for larger models.The paper presents this as an application of program transformations to inference.
A Iterative NUTS - Algorithm Details
The iterative NUTS algorithm replaces recursive tree construction with an iterative, memory-efficient procedure that integrates LeapFrog steps and checks U-Turn conditions.
- Recursive tree building: NUTS doubles trajectory length by running LeapFrog for twice the previous number of steps, building an implicit balanced binary tree.The process stops when a subtree trajectory begins to double back on itself.
- Recursive tree building: The recursive BuildTree procedure combines two depth d−1 subtrees while storing O(1) data per subtree for O(log N) memory.Here N = 2^d, avoiding storage of all N momentum-position pairs.
- Iterative tree building: The iterative BuildTree runs LeapFrog for N = 2^d steps from an initial position-momentum node and tree-depth argument.The procedure uses 0-based indexing.
- U-Turn condition: U-Turn checks compare each rightmost node with selected leftmost nodes from binary subtrees identified through progressively masked binary representations.Only odd-numbered nodes need checks against a subset of previous even-numbered nodes.
- Memory efficiency: A naive iterative implementation would require O(N) memory, so the algorithm stores only selected even-numbered nodes in an array indexed by bit count.The array can be overwritten during construction and has maximum size d.
- Algorithm comparison: Figure 4 contrasts recursive and iterative tree building at a high level, omitting proposal metadata, step size, and direction-selection details.Those omitted details are not relevant to the proposed changes.
B Code for Vectorized Sampling - Logistic Regression
The code defines a Pyro-compatible logistic regression model in NumPyro using JAX arrays, sampled coefficients and intercept, and Bernoulli observations.
- Model definition: The logistic regression model samples coefficient vector m and intercept b, then models y with Bernoulli logits x @ m + b.The implementation derives the number of dimensions from x and optionally conditions on observed y.
C Experimental Details
Experiments evaluate NumPyro and other frameworks across small-data HMM, large-scale logistic regression, and varying-dimensionality SKIM settings, using repeated runs and specified NUTS protocols.
- Experimental Setup: Experiments average results across 5 runs with different random seeds on an AMD Ryzen Threadripper 1920X and NVIDIA GeForce RTX 2080 Ti.The reported framework versions include PyStan 2.19.1.1, Pyro 1.0 with PyTorch 1.3.1, and NumPyro 0.2.3 with JAX 0.1.53 and jaxlib 0.1.36.
- Hidden Markov Model: The HMM benchmark uses 600 data points, 3-dimensional latent states, 10-dimensional observations, and 100 observed latent states.Stan and NumPyro use 1000 warmup steps and 1000 NUTS samples; Pyro uses a fixed 0.1 step size and 40 samples because its implementation is extremely slow.
- Logistic Regression: The Forest CoverType logistic-regression benchmark contains 581, 012 datapoints and 54 features after converting the multiclass task to binary classification.All frameworks use a fixed step size of 0.0015 and 40 samples.
- Sparse Kernel Interaction Model (SKIM): SKIM benchmarks use artificial datasets with N = 200 and 3 randomly selected pairwise interaction terms, varying the covariate dimensionality p.Each framework receives 1000 warmup adaptation steps, followed by 1000 samples for time-per-effective-sample estimates averaged over 5 random runs.
D Vectorized Estimation of the Evidence Lower Bound (ELBO) in SVI
The ELBO estimate is vectorized by splitting one random key into multiple particle keys and averaging mapped ELBO losses.
- Vectorized ELBO: VectorizedELBO computes the mean ELBO loss across num_particles random draws using random.split and vmap.The mapped losses call ELBO().loss with each split random key and the original arguments.