Simple, Distributed, and Accelerated Probabilistic Programming
Dustin Tran, Matthew W Hoffman, Dave Moore, Christopher Suter, Srinivas Vasudevan, Alexey Radul
Abstract
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.
BibTeX
@inproceedings{NEURIPS2018_201e5bac,
author = {Tran, Dustin and Hoffman, Matthew W and Moore, Dave and Suter, Christopher and Vasudevan, Srinivas and Radul, Alexey},
booktitle = {Advances in Neural Information Processing Systems},
editor = {S. Bengio and H. Wallach and H. Larochelle and K. Grauman and N. Cesa-Bianchi and R. Garnett},
pages = {},
publisher = {Curran Associates, Inc.},
title = {Simple, Distributed, and Accelerated Probabilistic Programming},
url = {https://proceedings.neurips.cc/paper_files/paper/2018/file/201e5bacd665709851b77148e225b332-Paper.pdf},
volume = {31},
year = {2018}
}