Source-linked AI summary
Pyro: Deep Universal Probabilistic Programming
Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, Noah D. Goodman
TL;DR
Advanced AI probabilistic models are often implemented as one-off systems, while existing probabilistic programming languages may require specialized inference for complex applications. Pyro addresses this by combining Python-based expressive modeling, PyTorch-backed stochastic variational inference, and composable mechanisms for model-specific behavior. Experiments report replication of a high-dimensional dynamical model and improved results from autoregressive-flow posteriors at negligible computational cost.
Problem
Advanced probabilistic models are often implemented as one-off systems, while practical probabilistic programming requires scalable inference for complex research models.
Method
Pyro combines Python-embedded probabilistic programs, PyTorch-backed stochastic variational inference, and Poutine’s composable effect handlers for flexible inference behavior.
Results
Pyro replicated the DMM’s reported model, inference configuration, and quantitative results, then improved results with autoregressive-flow posteriors at negligible computational cost.
Takeaways & Limitations
Pyro’s modular design supports concise replication and extensions of complex, high-dimensional probabilistic models.
Takeaways & Limitations
The paper notes that Pyro’s design decisions are not universally applicable or desirable, since other systems make different tradeoffs for different goals.
Abstract
from arXiv · showhide
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.
1. Introduction
Pyro addresses the gap between richly structured AI models and probabilistic programming systems by targeting scalable, flexible development for advanced research applications.
- Pyro targets probabilistic programming that can scale advanced AI research beyond one-off model implementations.The paper identifies one-off systems as slowing development and limiting scope and extensibility.
2. Design Principles
Pyro balances expressivity, scalability, flexibility, and minimality through Python embedding, scalable inference, and composable effect handlers. Its design choices address tensions among these objectives while supporting complex models and custom inference behavior.
- Expressivity: Expressivity requires support for data-dependent control flow, dependent latent-variable existence, and unnormalized joint distributions.
- Scalability: Scalability requires approximate inference for large datasets and non-conjugate, high-dimensional models, with compiler acceleration when possible.
- Flexibility: Flexibility supports complex model-specific inference behavior while separating model, inference, and runtime implementations.
- Minimality: Minimality reduces cognitive overhead by sharing syntax and semantics with existing systems and interoperating with visualization libraries.
- Pyro’s design: Pyro balances conflicting design objectives by embedding probabilistic programs in Python and exposing two language primitives, pyro.sample and pyro.param.Python control flow, including recursion, loops, and conditionals, can determine which random variables exist during execution.
- Inference and runtime: Pyro’s primary inference algorithm is gradient-based stochastic variational inference, while its tensor math and automatic differentiation rely on PyTorch.Poutine separates inference algorithm implementations from language details through composable effect handlers.
3. Project Openness and Development
Pyro is developed as an open-source project with public code, documentation, examples, and testing. Its development also addressed a missing high-performance probability-distribution library in PyTorch through upstream contributions.
- Pyro’s source code is available under an MIT license, with documentation, examples, and a discussion forum hosted online.
- Pyro uses continuous integration to run a comprehensive test suite before code is merged, supporting project quality and usability.
- PyTorch lacked a high-performance probability-distribution library, prompting substantial contributions to PyTorch Distributions.
4. Existing Systems
The paper situates Pyro among probabilistic programming systems that make different tradeoffs across expressivity, scalability, flexibility, and minimality. It emphasizes that Pyro’s design choices are not universally applicable or desirable.
- Existing probabilistic programming systems differ in host language, control flow, subsampling and automatic differentiation, flexible inference, and tool integration.
- The paper presents Figure 2 as a simplified summary of design principles for Pyro and other probabilistic programming languages.
- Pyro’s design decisions are explicitly not claimed to be universally applicable, because other systems purposefully make different tradeoffs for different goals.
- The comparison includes Stan, Church, Venture, Anglican, webPPL, Edward, ProbTorch, and Turing, which target distinct combinations of expressivity, inference, and implementation goals.
5. Experiments
Pyro’s experiments evaluate state-of-the-art VAE and DMM models, testing scalability against PyTorch and flexibility through richer DMM guides. The implementations reproduce reference results while enabling expressive extensions with low computational cost.
- Experimental models: Pyro evaluates a VAE and the high-dimensional, non-conjugate DMM on large-data settings with data-dependent latent variables.The DMM also uses a hand-designed approximate posterior, making it a point of comparison for Pyro’s flexibility and expressiveness.
- Scalability: The Pyro and PyTorch VAE implementations converge to the same test ELBO, with a moderate relative performance gap that shrinks as tensor-operation time increases.Timing uses GPU-accelerated mini-batch stochastic gradient variational inference with batch size 128 on a single NVIDIA GTX 1080Ti.
- Flexibility and expressiveness: Pyro reproduces the DMM model, inference configuration, and quantitative results after 5000 training epochs.The replication is described as exact and concise relative to the reported DMM paper.
- Flexibility and expressiveness: IAF-based DMM variants improve results at negligible computational cost and require only a few lines of code.The extension uses more expressive approximate posteriors enabled by Pyro’s modular design.