# Simple, Distributed, and Accelerated Probabilistic Programming

arXiv: https://arxiv.org/abs/1811.02091 · HTML: https://mlsanity.com/paper/1811.02091/

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 5 Nov 2018.

## Built on by (most cited first, 18 of 18)

- [Plex: Towards Reliability using Pretrained Large Model Extensions](https://mlsanity.com/paper/2207.07411/) (2022)
- [Tensor Program Optimization with Probabilistic Programs](https://mlsanity.com/paper/2205.13603/) (2022)
- [RecSim NG: Toward Principled Uncertainty Modeling for Recommender Ecosystems](https://mlsanity.com/paper/2103.08057/) (2021)
- [Learning Proposals for Probabilistic Programs with Inference Combinators](https://mlsanity.com/paper/2103.00668/) (2021)
- [TyXe: Pyro-based Bayesian neural nets for Pytorch](https://mlsanity.com/paper/2110.00276/) (2021)
- [Smoothness Analysis for Probabilistic Programs with Application to Optimised Variational Inference](https://mlsanity.com/paper/2208.10530/) (2022)
- [A Survey on Uncertainty Toolkits for Deep Learning](https://mlsanity.com/paper/2205.01040/) (2022)
- [Probabilistic Programming with Programmable Variational Inference](https://mlsanity.com/paper/2406.15742/) (2024)
- [D3p -- A Python Package for Differentially-Private Probabilistic Programming](https://mlsanity.com/paper/2103.11648/) (2021)
- [A Heavy-Tailed Algebra for Probabilistic Programming](https://mlsanity.com/paper/2306.09262/) (2023)
- [Distilled Thompson Sampling: Practical and Efficient Thompson Sampling via Imitation Learning](https://mlsanity.com/paper/2011.14266/) (2020)
- [Meta-Learning an Inference Algorithm for Probabilistic Programs](https://mlsanity.com/paper/2103.00737/) (2021)
- [Higher-Order Generalization Bounds: Learning Deep Probabilistic Programs via PAC-Bayes Objectives](https://mlsanity.com/paper/2203.15972/) (2022)
- [Variational Inference for Category Recommendation in E-Commerce platforms](https://mlsanity.com/paper/2104.07748/) (2021)
- [BayesLDM: A Domain-Specific Language for Probabilistic Modeling of Longitudinal Data](https://mlsanity.com/paper/2209.05581/) (2022)
- [Push: Concurrent Probabilistic Programming for Bayesian Deep Learning](https://mlsanity.com/paper/2306.06528/) (2023)
- [YAMLE: Yet Another Machine Learning Environment](https://mlsanity.com/paper/2402.06268/) (2024)
- [Representing Molecules with Algebraic Data Types: Beyond SMILES and SELFIES](https://mlsanity.com/paper/2501.13633/) (2025)

## Abstract (arXiv)

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.