Pyro: Deep Universal Probabilistic Programming

Eli BinghamJonathan P. ChenMartin JankowiakFritz ObermeyerNeeraj PradhanTheofanis KaraletsosRohit SinghPaul SzerlipPaul HorsfallNoah D. Goodman

article2018JMLR1,337 citations

Presents Pyro, a universal probabilistic programming language built on PyTorch that couples deep neural networks with stochastic variational inference to scale expressive Bayesian models to high-dimensional datasets.

Listen

Developing advanced artificial intelligence applications increasingly relies on complex probabilistic models. Historically, implementing these architectures required engineering custom, one-off systems from scratch, which consumed substantial development time and limited project scalability. Probabilistic programming languages aim to automate this development, but existing systems have struggled to simultaneously deliver high performance on large datasets, support dynamic model structures, and offer customizable inference algorithms.

The article demonstrates the design, implementation, and empirical validation of Pyro, an open-source probabilistic programming language embedded in Python and built on PyTorch. The primary objective is to evaluate whether Pyro successfully reconciles the trade-offs between expressive modeling, computational scalability, inference flexibility, and minimal syntax for modern machine learning workflows.

To evaluate the system, the authors designed Pyro around two core language primitives and leveraged PyTorch for graphics processing unit acceleration, reverse-mode automatic differentiation, and mini-batch stochastic variational inference. They validated the platform through controlled computational benchmarks using standard deep learning architectures, including a variational autoencoder on image data and a complex Deep Markov Model on sequential music data. These experiments measured per-update runtime overhead against custom PyTorch code and assessed model accuracy across various inference configurations.

The article presents three core findings. First, Pyro achieves competitive computational efficiency: in variational autoencoder benchmarks, Pyro introduced a moderate runtime overhead per update compared to hand-written PyTorch code (around 6.7 milliseconds versus 3.8 milliseconds on smaller networks), but this performance gap narrowed significantly (to roughly 10.2 milliseconds versus 7.7 milliseconds) as network size and compute workload grew. Second, the system fully replicated established baseline results for high-dimensional models with dynamic sequence lengths, matching prior Deep Markov Model performance. Third, Pyro demonstrated rapid extensibility, allowing researchers to integrate inverse autoregressive flows into the Deep Markov Model inference pipeline with only a few lines of code, improving quantitative lower-bound test scores with negligible computational overhead.

These results show that organizations can significantly reduce development cycles and engineering costs for complex artificial intelligence research without sacrificing execution speed on large-scale datasets. By decoupling model definitions from inference mechanics, teams can lower project risk, accelerate prototyping timelines, and easily modify state-of-the-art architectures without rebuilding low-level infrastructure from scratch.

Technical leaders looking to implement deep probabilistic models should consider adopting Pyro for workflows requiring dynamic control flows and custom approximate inference. Teams planning adoption should use GPU acceleration and mini-batching to minimize software overhead. Furthermore, organizations investing in the platform can leverage and contribute to its open-source ecosystem, particularly upstream libraries supporting probability distributions.

The article evaluates performance across a focused set of benchmark models and datasets on a single hardware configuration. While confidence is high that Pyro efficiently handles the tested deep generative architectures, teams with non-standard hardware setups, specialized distributed computing requirements, or applications requiring non-variational inference algorithms should conduct their own internal pilot benchmarks before wide-scale deployment.

Stanford University Uber deep learning gpu inference deep universal probabilistic programming probabilistic programming languages Pyro

arXiv: 1810.09538

Paper: Automatic differentiation in PyTorch, Adam Paszke et al. (2017). Pyro is directly built on top of PyTorch's imperative automatic differentiation engine to execute and differentiate probabilistic programs dynamically.
Paper: Stochastic Backpropagation and Approximate Inference in Deep Generative Models, Danilo Jimenez Rezende et al. (2014). It establishes the foundational stochastic gradient variational Bayes and reparameterization techniques that Pyro implements for deep universal probabilistic inference.
Paper: Tutorial on Variational Autoencoders, Carl Doersch (2016). It provides a clear theoretical foundation for variational autoencoders and stochastic variational inference that underpin model-guide pairs in Pyro.
Paper: Variational Inference with Normalizing Flows, Danilo Jimenez Rezende et al. (2015). It introduces normalizing flows to enrich approximate posteriors in variational inference, a core family of distributions supported within Pyro.
Paper: Improving Variational Inference with Inverse Autoregressive Flow, Diederik P. Kingma et al. (2016). It develops inverse autoregressive flows, providing scalable, expressive variational distributions that can be seamlessly formulated in deep probabilistic programming languages.
Paper: The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables, Chris J. Maddison et al. (2016). It introduces continuous relaxations for discrete latent variables, enabling low-variance gradient-based variational inference in probabilistic programs.
Paper: Weight Uncertainty in Neural Network, Charles Blundell et al. (2015). It details Bayes by Backprop for estimating weight uncertainty in neural networks, representing a key use case for universal probabilistic programming languages like Pyro.

Paper: PyTorch: An Imperative Style, High-Performance Deep Learning Library, Adam Paszke et al. (2019). This paper provides the complete architectural design and benchmarked ecosystem of PyTorch, the underlying framework that powers Pyro's tensor operations and autograd runtime.
Paper: Normalizing Flows for Probabilistic Modeling and Inference, George Papamakarios et al. (2019). This survey provides an extensive, formalized analysis of normalizing flows for probabilistic modeling and inference, extending the advanced distribution components used in Pyro.
Paper: The frontier of simulation-based inference, Kyle Cranmer et al. (2019). It surveys simulation-based inference methods that generalize and build upon probabilistic programming systems and neural surrogates for implicit simulator models.
Paper: An Introduction to Variational Autoencoders, Diederik P. Kingma et al. (2019). It expands upon variational autoencoders and amortized variational inference frameworks that Pyro is designed to express and scale.
Paper: Decision-Making with Auto-Encoding Variational Bayes, Romain Lopez et al. (2020). It builds upon amortized variational inference models to improve downstream decision-making and mitigate bias in estimated posterior expectations.
Paper: PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation, Jason Ansel et al. (2024). It details graph compilation and dynamic bytecode transformations that modernize the PyTorch execution backend utilized by downstream probabilistic programming systems.

Cover for Pyro: Deep Universal Probabilistic Programming

Abstract

Pyro is a probabilistic programming language built on Python as a platform for developing advanced probabilistic models in AI research. To scale to large datasets and high-dimensional models, Pyro uses stochastic variational inference algorithms and probability distributions built on top of PyTorch, a modern GPU-accelerated deep learning framework. To accommodate complex or model-specific algorithmic behavior, Pyro leverages Poutine, a library of composable building blocks for modifying the behavior of probabilistic programs.

Topics

deep learning gpu inference inference algorithm libraries probabilistic models programming languages python variational inference variational approximation

Keywords

deep universal probabilistic programming probabilistic programming languages Pyro deep generative models stochastic variational inference Hamiltonian Monte Carlo variational autoencoder automatic differentiation approximate inference

1 Introduction
2 Design Principles
3 Project Openness and Development
4 Existing Systems
5 Experiments
References

Knowls

Knowl 1 — Core Language Primitives and Execution Model of Pyro
model/method
Pyro is a universal probabilistic programming language embedded in Python that represents probabilistic models as standard Python functions or callables. Models support dynamic Python control flow, including loops, recursion, and data-dependent branching, allowing the collection of latent random variables in an execution trace to vary dynamically.

Pyro defines two core primitives to interact with inference engines:
1. pyro.sample(name, distribution, obs=None): Annotates a call to a random distribution with an identifier name. If obs is omitted, the primitive defines an unobserved latent random variable sampled from distribution. If obs=val is provided (or when the program is modified via conditioning), the statement binds the random variable to observed data val for computing unnormalized joint probabilities.
2. pyro.param(name, init_tensor): Registers a named learnable deterministic parameter with Pyro's central parameter store. Inference algorithms optimize these parameters across successive executions.
Interoperability with neural networks is achieved via pyro.module, which registers all parameters of a torch.nn.Module with the parameter store under pyro.param.
Knowl 2 — Programmable Variational Posteriors (Guides) in Pyro
model/method
In Pyro, approximate posterior distributions and proposal distributions for stochastic variational inference (SVI) and importance sampling are defined as explicit, standalone probabilistic programs termed guides.

A guide program must satisfy the following constraints with respect to its target generative model:
1. Input signature compatibility: The guide must accept the exact same input arguments as the model.
2. Sample site correspondence: The guide must contain a corresponding pyro.sample statement for every unconstrained (unobserved) pyro.sample statement in the model, using the same string identifier.
3. Observation exclusion: The guide must not contain sample statements corresponding to observed data sites in the model.
4. Scope separation: Unlike languages like webPPL or Anglican, a Pyro guide cannot read or condition on values sampled within the model trace during inference; dependencies must flow through input arguments or shared deterministic parameters registered in the parameter store.
This abstraction allows the user to specify non-mean-field dependencies, deep amortized neural network encoders, or normalising flow transformations in the posterior.
Knowl 3 — Poutine Effect Handler Architecture in Pyro
model/method
Pyro achieves modularity and separation of concerns between model specification and inference algorithms through Poutine, an internal library of algebraic effect handlers. Poutine intercepts, records, and modifies the execution of probabilistic primitives (pyro.sample and pyro.param) at runtime.

Key effect handling behaviors implemented in Poutine include:
- Trace recording: Intercepting primitive calls to construct an execution trace graph recording sample values, log-joint probabilities, parameter references, and computational dependencies.
- Conditioning (pyro.condition): Dynamically replacing generative sampling statements with fixed observed values without altering the original model code.
- Replay: Forcing model execution to reuse the exact latent variable values sampled during a corresponding guide execution trace.
- Parameter interception: Managing lookup, scoping, and parameter retrieval from the central parameter store.
By composing these primitive effect handlers, inference algorithms operate purely as transformations over execution traces rather than requiring specialized compiler passes or model modifications.
Knowl 4 — Stochastic Variational Inference via Trace ELBO
algorithm
Stochastic Variational Inference (SVI) in Pyro optimizes parameters $\theta$ of a model $p_\theta(x, z)$ and parameters $\phi$ of a guide $q_\phi(z)$ by maximizing the Evidence Lower Bound (ELBO):

$\text{ELBO}(\theta, \phi) = \mathbb{E}_{q_\phi(z)} \left[ \log p_\theta(x, z) - \log q_\phi(z) \right]$

where $x$ denotes observed data and $z$ denotes latent variables. Rather than requiring analytic expectations, Pyro computes single-sample or multi-sample Monte Carlo estimates of the log-joint probability ratio across execution traces.
Input: Generative model $p_\theta(x, z)$ , guide program $q_\phi(z)$ , dataset $D$ , learning rate $\alpha$ , mini-batch size $B$
Output: Optimized parameter tensors $\theta$ and $\phi$
Initialize parameters $\theta$ and $\phi$ in the Pyro parameter store
for each training iteration until convergence do
    Sample a mini-batch of data $\{x_i\}_{i=1}^B$ from $D$
    for each data point $x_i$ in mini-batch do
        Execute guide $q_\phi(z \mid x_i)$ to generate latent trace $z^{(i)} \sim q_\phi(z \mid x_i)$
        Compute guide log-probability $\log q_\phi(z^{(i)} \mid x_i)$
        Replay model $p_\theta(x_i, z)$ conditioned on $z^{(i)}$ to evaluate joint log-probability $\log p_\theta(x_i, z^{(i)})$
        Compute Monte Carlo ELBO estimate: $\widehat{\text{ELBO}}_i = \log p_\theta(x_i, z^{(i)}) - \log q_\phi(z^{(i)} \mid x_i)$
    Compute mini-batch loss: $\mathcal{L}(\theta, \phi) = -\frac{1}{B} \sum_{i=1}^B \widehat{\text{ELBO}}_i$
    Compute parameter gradients $\nabla_\theta \mathcal{L}$ and $\nabla_\phi \mathcal{L}$ via reverse-mode automatic differentiation
    Update parameters: $(\theta, \phi) \leftarrow (\theta, \phi) - \alpha \cdot \text{Adam}(\nabla_{\theta, \phi} \mathcal{L})$
return optimized parameter tensors $\theta, \phi$

Knowl 5 — Design Principles Comparison Across Probabilistic Programming Languages

data/table

Probabilistic programming languages (PPLs) are differentiated along four design dimensions: expressivity (support for dynamic control flow versus static computation graphs), scalability (native support for dataset subsampling and hardware-accelerated automatic differentiation), flexibility (support for programmable, user-defined inference), and minimality (embedding within host languages).

System	Dynamic Control Flow	Subsampling, AD	Flexible Inference	Host Language
Stan	Static control flow	Some, CPU	Automated	None
Church	Yes	No, None	Automated	Scheme
Venture	Yes	No, None	Yes	None
webPPL	Yes	No, CPU	Some	JavaScript
Edward	Static control flow	Yes, CPU/GPU	Yes	TensorFlow
Pyro	Yes	Yes, CPU/GPU	Yes	Python, PyTorch

Pyro combines full dynamic control flow (expressiveness) with GPU-accelerated automatic differentiation and mini-batch stochastic variational inference (scalability) while embedding into standard Python.

Knowl 6 — Per-Update Execution Time Comparison Between Pyro and PyTorch on VAE

data/table

To evaluate the computational overhead introduced by Pyro's effect handlers and dynamic trace recording, a Variational Autoencoder (VAE) trained on MNIST with 2-hidden-layer MLP encoders and decoders was benchmarked against an idiomatic raw PyTorch implementation. The benchmark reports wall-clock time per mini-batch gradient update (milliseconds) averaged over 10 epochs on an NVIDIA GTX 1080Ti GPU (batch size 128) across latent dimension size ( $\#z$ ) and hidden layer dimension ( $\#h$ ).

$\#z$ (Latent Dim)	$\#h$ (Hidden Dim)	PyTorch (ms)	Pyro (ms)
10	400	$3.82 \pm 0.02$	$6.79 \pm 0.04$
30	400	$3.73 \pm 0.07$	$6.67 \pm 0.03$
10	2000	$7.65 \pm 0.02$	$10.14 \pm 0.06$
30	2000	$7.66 \pm 0.02$	$10.19 \pm 0.03$

Both implementations converge to identical test ELBO values. Pyro incurs an absolute overhead of approximately $2.5$ to $3.0\text{ ms}$ per step due to trace management. As tensor dimensions increase (from $\#h=400$ to $\#h=2000$ ), tensor operations dominate total computation time, reducing the relative overhead from approximately $78\%$ to $33\%$ .

Knowl 7 — Deep Markov Model Performance with Inverse Autoregressive Flow Guides

empirical result

A Deep Markov Model (DMM)—a non-linear state space generative model where the number of latent variables scales dynamically with sequence length—was implemented in Pyro and evaluated on the JSB Chorales dataset over 5,000 training epochs. Pyro replicated the baseline test Evidence Lower Bound (ELBO) and supported extending the approximate posterior guide with Inverse Autoregressive Flows (IAFs) with minimal additional code:

Augmenting the DMM guide with one and two IAF normalizing flow transformations improved the test ELBO from $-6.87$ to $-6.82$ and $-6.80$ respectively, illustrating Pyro's capacity to incorporate expressive structured inference networks into dynamic models.

Coverage note — General software engineering details (open-source MIT licensing, repository links, CI testing setup, and upstream contributions to the PyTorch Distributions library) were omitted as they do not constitute standalone scientific contributions.

References

1.Bob Carpenter, Andrew Gelman, Matthew D. Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A Probabilistic Programming Language. Journal of Statistical Software, 76(1), 2017.
2.Joshua V. Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A. Saurous. TensorFlow Distributions. arXiv:1711.10604, November 2017.
3.Hong Ge, Kai Xu, and Zoubin Ghahramani. Turing: A Language for Flexible Probabilistic Inference. In AISTATS, 2018.
4.Zoubin Ghahramani. Probabilistic machine learning and artificial intelligence. Nature, 521: 452–459, May 2015.
5.Noah D Goodman and Andreas Stuhlmüller. The Design and Implementation of Probabilistic Programming Languages. http://dippl.org, 2014.
6.Noah D. Goodman, Vikash K. Mansinghka, Daniel Roy, Keith Bonawitz, and Joshua B. Tenenbaum. Church: A Language for Generative Models. In UAI, 2008.
7.Matthew D. Hoffman and Andrew Gelman. The No-U-turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo. J. Mach. Learn. Res., 15(1), January 2014.
8.Ohad Kammar, Sam Lindley, and Nicolas Oury. Handlers in Action. In ICFP, 2013.
9.Diederik P Kingma and Max Welling. Auto-encoding Variational Bayes. In ICLR, 2014.
10.Diederik P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved Variational Inference with Inverse Autoregressive Flow. In NIPS. 2016.
11.Rahul G Krishnan, Uri Shalit, and David Sontag. Structured Inference Networks for Nonlinear State Space Models. In AAAI, 2017.
12.Vikash K. Mansinghka, Ulrich Schaechtle, Shivam Handa, Alexey Radul, Yutian Chen, and Martin Rinard. Probabilistic Programming with Programmable Inference. In PLDI, 2018.
13.N. Siddharth, Brooks Paige, Jan-Willem van de Meent, Alban Desmaison, Noah D. Goodman, Pushmeet Kohli, Frank Wood, and Philip Torr. Learning Disentangled Representations with Semi-Supervised Deep Generative Models. In NIPS, 2017.
14.David Tolpin, Jan-Willem van de Meent, Hongseok Yang, and Frank Wood. Design and Implementation of Probabilistic Programming Language Anglican. In IFL, 2016.
15.Dustin Tran, Matthew D. Hoffman, Rif A. Saurous, Eugene Brevdo, Kevin Murphy, and David M. Blei. Deep Probabilistic Programming. In ICLR, 2017.
16.David Wingate, Andreas Stuhlmller, and Noah Goodman. Lightweight Implementations of Probabilistic Programming Languages via Transformational Compilation. In AISTATS, 2011.

Citation

MLA

Bingham, E., et al. “Pyro: Deep Universal Probabilistic Programming”. arXiv, 2018, http://arxiv.org/abs/1810.09538v1.

APA

Bingham, E., Chen, J. P., Jankowiak, M., Obermeyer, F., Pradhan, N., Karaletsos, T., Singh, R., Szerlip, P., Horsfall, P., & Goodman, N. D. (2018). Pyro: Deep Universal Probabilistic Programming. arXiv. http://arxiv.org/abs/1810.09538v1

Chicago

Bingham, E., J. P. Chen, M. Jankowiak, et al. 2018. “Pyro: Deep Universal Probabilistic Programming”. arXiv. http://arxiv.org/abs/1810.09538v1.

Harvard

Bingham, E. et al. (2018) “Pyro: Deep Universal Probabilistic Programming”, arXiv [Preprint]. Available at: http://arxiv.org/abs/1810.09538v1.

Vancouver

1. Bingham E, Chen JP, Jankowiak M, Obermeyer F, Pradhan N, Karaletsos T, Singh R, Szerlip P, Horsfall P, Goodman ND (2018) Pyro: Deep Universal Probabilistic Programming. arXiv

BibTeX

@article{bingham2018pyro,
  title = {Pyro: Deep Universal Probabilistic Programming},
  author = {Bingham, Eli and Chen, Jonathan P. and Jankowiak, Martin and Obermeyer, Fritz and Pradhan, Neeraj and Karaletsos, Theofanis and Singh, Rohit and Szerlip, Paul and Horsfall, Paul and Goodman, Noah D.},
  year = {2018},
  journal = {arXiv},
  url = {http://arxiv.org/abs/1810.09538v1},
  eprint = {1810.09538}
}

Metadata:arXiv

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

Guide Configuration	Test ELBO
0 IAFs (Krishnan et al., 2017 baseline)	-6.93
0 IAFs (Pyro replication)	-6.87
Pyro with 1 IAF	-6.82
Pyro with 2 IAFs	-6.80