Simple, Distributed, and Accelerated Probabilistic Programming

By Dustin Tran, Matthew D. Hoffman, David A. Moore, Christopher Suter, Srinivas Vasudevan, Alexey Radul, Matthew Johnson, Rif A. Saurous. 2018

61 citations recorded by the source corpus.

We describe a simple, low-level approach for embedding probabilistic programming in a deep learning ecosystem. In particular, we distill probabilistic programming down to a single abstraction---the random variable. Our lightweight implementation in TensorFlow enables numerous applications: a model-parallel variational auto-encoder (VAE) with 2nd-generation tensor processing units (TPUv2s); a data-parallel autoregressive model (Image Transformer) with TPUv2s; and multi-GPU No-U-Turn Sampler (NUTS). For both a state-of-the-art VAE on 64x64 ImageNet and Image Transformer on 256x256 CelebA-HQ, our approach achieves an optimal linear speedup from 1 to 256 TPUv2 chips. With NUTS, we see a 100x speedup on GPUs over Stan and 37x over PyMC3.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, 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. Gilles Louppe
2. Kyle Cranmer
3. Yura N. Perov
4. Hongseok Yang
5. Frank D. Wood
6. Danijar Hafner
7. Dustin Tran
8. Atilim Gunes Baydin

In Computer Vision and Pattern Recognition

1. Gilles Louppe
2. Kyle Cranmer
3. Michael W. Dusenberry
4. Danijar Hafner
5. Mark van der Wilk
6. Frank D. Wood
7. Gwonsoo Che
8. Atilim Gunes Baydin

In Probabilistic & Bayesian Methods

1. Gilles Louppe
2. Kyle Cranmer
3. Hongseok Yang
4. Frank D. Wood
5. Yura N. Perov
6. Atilim Gunes Baydin
7. Adam Baker
8. Saurabh Johri

Citations in the corpus

• Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro (2019, methodology): “Figure 3 illustrates the iterative procedure at step n = 11 , and full details of the algorithm are provided in Algorithm 2.”
• Bayesian Layers: A Module for Neural Network Uncertainty (2018, methodology): “Dependency-wise, it extends Keras in TensorFlow (Chollet, 2016) and uses Edward2 (Tran et al., 2018) to operate with random variables.”
• Efficient Probabilistic Inference in the Quest for Physics Beyond the Standard Model (2018, methodology): “[83] D.”
• MultiVerse: Causal Reasoning using Importance Sampling in Probabilistic Programming (2019, methodology): “Counterfactual Inference in Existing Probabilistic Programming Frameworks and Comparison to MultiVerse Given an intervention mechanism such as exists natively in Pyro (Bingham et al., 2018) (or as can be implemented in Edward as in Tran et al. (2018)), one can write the steps of abduction, intervent”
• Meta-Learning an Inference Algorithm for Probabilistic Programs (2021, background): “…specifying conditional independence in a model (Bingham et al., 2018) or defining proposal or variational distributions (Ritchie et al., 2015; Siddharth et al., 2017; Bingham et al., 2018; Tran et al., 2018; Cusumano-Towner et al., 2019), all enabling users to help inference or learning algorithms.”