Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro

By Du Phan, Neeraj Pradhan, Martin Jankowiak. 2019

531 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computational Theory and Mathematics, Hardware and Architecture, Probabilistic & Bayesian Methods

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Andrew Gordon Wilson
2. Quoc V. Le
3. Yoshua Bengio
4. Wesley J. Maddox
5. Yuhuai Wu
6. Sergey Levine
7. Jacob Austin
8. Jascha Narain Sohl-Dickstein

In Computational Theory and Mathematics

1. Jascha Narain Sohl-Dickstein
2. Osvaldo A. Martin
3. Andrew Gordon Wilson
4. Corey Nolet
5. Wesley J. Maddox
6. Martin Hirzel
7. Gregory W. Benton
8. Joshua Patterson

In Hardware and Architecture

1. Andrew Gordon Wilson
2. Corey Nolet
3. Wesley J. Maddox
4. Gregory W. Benton
5. Joshua Patterson
6. Sebastian Raschka
7. Thomas V. Wiecki
8. Yuhuai Wu

In Probabilistic & Bayesian Methods

1. Andrew Gordon Wilson
2. Wesley J. Maddox
3. Yuhuai Wu
4. Jacob Austin
5. Louis Mandel
6. Guillaume Baudart
7. Corey Nolet
8. Joshua Patterson

Citations in the corpus

• Language Model Cascades (2022, methodology): “We implement cascades as a trace-based probabilistic programming language embedded in Python via effect handlers, inspired by Bingham et al. (2018); Phan et al. (2019), and via coroutines, inspired by Kochurov et al. (2019).”
• Rethinking Parameter Counting in Deep Models: Effective Dimensionality Revisited (2020, methodology): “4 We then construct a Bayesian neural network with two hidden layers each with 20 units, no biases, and tanh activations, placing independent Gaussian (4)From the Bayesian neural network example in NumPyro (Phan et al., 2019; Bingham et al., 2019): https://github.”
• PyMC: a modern, and comprehensive probabilistic programming framework in Python (2023, methodology): “Currently, PyMC offers NUTS JAX samplers via NumPyro ( Phan, Pradhan & Jankowiak, 2019 ) or BlackJAX ( BlackJax devs, 2022 ) with the functions pm.sample_numpyro_nuts and pm. sample_blackjax_nuts , respectively.”
• Machine Learning in Python: Main developments and technology trends in data science, machine learning, and artificial intelligence (2020, methodology): “NumPyro [228] provides a NumPy backend for Pyro, using JAX to JIT- compile and optimize execution of NumPy operations on both CPUs and GPUs.”
• Towards Understanding Sycophancy in Language Models (2023, unknown): “These results suggest the Claude 2 PM sometimes prefers sycophantic responses over truthful ones. was trained on a mix of human preference judgments and AI preference judgments (Anthropic, 2023).”
• Bean Machine: A Declarative Probabilistic Programming Language For Efficient Programmable Inference (2020, methodology): “, 2017) and NumPyro (Phan et al., 2019) are all using NUTS.”
• Compiling Stan to generative probabilistic languages and extension to deep probabilistic programming (2021, methodology): “4 Implementation We implemented two new backends for the Stan compiler targeting Pyro [3] and NumPyro [25].”