[Paper Review] Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro

Du Phan, Neeraj Pradhan+1 more|arXiv (Cornell University)|Dec 24, 2019

Parallel Computing and Optimization Techniques15 references226 citations

TL;DR

The paper introduces NumPyro, a NumPy-based probabilistic programming library that uses composable effect handlers with JAX transformations to achieve end-to-end JIT compilation and significant speedups, including an iterative NUTS sampler that outperforms existing implementations.

ABSTRACT

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.

Motivation & Objective

• Demonstrate that Pyro-like effect handlers can be integrated with a JAX-based backend to extend NumPyro’s modeling and inference interface.
• Show how composable transformations (jit, grad, vmap) accelerate inference subroutines while preserving a Pyro-compatible modeling language.
• Present an iterative, end-to-end JIT-compiled NUTS implementation that exploits XLA for speedups on CPU and GPU.
• Illustrate the benefits of vectorizing subroutines with vmap to enable batched inference and predictions.
• Benchmark NumPyro against Stan and Pyro to quantify performance improvements across small and large datasets.

Proposed method

• Adopt a Pyro-like effect-handling abstraction (seed, trace, condition) within NumPyro to enable composable program transformations on a JAX backend.
• Leverage JAX transformations (jit, grad, vmap) to build and optimize inference subroutines with functional, traceable code.
• Transform the NUTS BuildTree subroutine into an iterative form to enable end-to-end JIT compilation and better operator fusion via XLA.
• Use vmap to batch common inference tasks such as prior sampling, posterior predictive sampling, and log-likelihood computation.
• Demonstrate the integration of effect handlers with JAX transformations to maintain a Pyro-compatible modeling interface while gaining acceleration.

Experimental results

Research questions

• RQ1 Can Pyro-style effect handlers be effectively composed with JAX transformations to accelerate probabilistic programming workloads in NumPyro?
• RQ2 What performance gains arise from end-to-end JIT compilation of inference subroutines like NUTS in NumPyro compared to Pyro and Stan?
• RQ3 How does vectorization via vmap interact with effect handlers to enable scalable batched inference and predictions?
• RQ4 Does an iterative, JIT-compiled NUTS implementation maintain correctness while delivering substantial speedups across datasets of varying sizes?
• RQ5 What are the practical gains in CPU, GPU, and TPU backends when using NumPyro’s composable transformations for large-scale models?

Key findings

• NumPyro’s iterative NUTS implementation is significantly faster than Pyro and Stan for the tested models, with reported speedups of around 340x over Pyro and 6x over Stan on certain tasks.
• End-to-end JIT compilation via XLA yields substantial speedups by enabling full fusion and optimized execution of inference subroutines.
• Vectorizing inference with vmap enables efficient batched sampling from the prior and posterior predictive distributions as well as batched log-likelihood computations.
• NumPyro’s framework maintains a Pyro-compatible modeling language while leveraging JAX transformations for acceleration across CPU, GPU, and TPU backends.
• The paper shows end-to-end acceleration benefits are observed in both small-scale CPU regimes and large-scale GPU regimes, motivating broader applicability of composable effects with JAX.